[BUG] module 'optree' has no attribute 'PyTreeDef'

woithook commented 1 year ago

Describe the bug

I used torchopt in google colab without any problem untill last week, an error occurred when I tried to import torchopt.

AttributeError: module 'optree' has no attribute 'PyTreeDef'

I have tried reinstalling torchop as well as pytorch, and even replacing a local machine to execute the program, but the problem has not been solved.

To Reproduce

The error occurs when executing the following code on Google Colab:

!pip3 install torchopt
import torchopt

Then

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In [2], line 1
----> 1 import torchopt

File ~/.pyenv/versions/anaconda3-5.3.1/envs/meta_rl/lib/python3.8/site-packages/torchopt/__init__.py:17
      1 # Copyright 2022 MetaOPT Team. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     13 # limitations under the License.
     14 # ==============================================================================
     15 """TorchOpt: a high-performance optimizer library built upon PyTorch."""
---> 17 from torchopt._src import (
     18     accelerated_op_available,
     19     clip,
     20     combine,
     21     hook,
     22     implicit_diff,
     23     linear_solve,
     24     schedule,
     25     visual,
     26 )
     27 from torchopt._src.alias import adam, adamw, rmsprop, sgd
     28 from torchopt._src.clip import clip_grad_norm

File ~/.pyenv/versions/anaconda3-5.3.1/envs/meta_rl/lib/python3.8/site-packages/torchopt/_src/implicit_diff.py:159
    153     out_kwargs = {k: v for k, (_, _, v) in bound.kwargs.items()}
    154     return out_args, out_kwargs, map_args_back
    157 def _split_tensor_and_others(
    158     mixed_tuple: Tuple,
--> 159 ) -> Tuple[pytree.PyTreeDef, Tuple[bool, ...], Tuple[torch.Tensor, ...], Tuple[Any, ...]]:
    160     flattened, treedef = pytree.tree_flatten(mixed_tuple)
    161     tensors = []

AttributeError: module 'optree' has no attribute 'PyTreeDef'

System info

CUDA Version: 11.0
python 3.8
torch 1.13.0

Checklist

[x] I have checked that there is no similar issue in the repo. (required)
[x] I have read the documentation. (required)
[x] I have provided a minimal working example to reproduce the bug. (required)

XuehaiPan commented 1 year ago

@woithook Thanks for raising this. We accidentally removed the type alias in optree. You can install the older version as:

pip3 install optree==0.3.0

We will read the alias and make a new release of optree soon.

XuehaiPan commented 1 year ago

We released optree 0.5.0 to PyPI, which readded the type alias.

metaopt / torchopt