Repository: rainx0r/metaworld-algorithms
Issue: AttributeError: module 'jax.experimental.layout' has no attribute 'DeviceLocalLayout'
Root cause: External Dependency Breakage

I followed the instructions in the readme file to set up the environment on Ubuntu 22.04, but encountered this issue when running. Could you provide me with the correct version of JAX?
```
Traceback (most recent call last):
  File "/home/rm/wxz/EmbodiedAI/metaworld/metaworld-algorithms/examples/multi_task/moore_mt10.py", line 14, in <module>
    from metaworld_algorithms.rl.algorithms import MTSACConfig
  File "/home/rm/wxz/EmbodiedAI/metaworld/metaworld-algorithms/metaworld_algorithms/rl/algorithms/__init__.py", line 3, in <module>
    from .base import Algorithm, OffPolicyAlgorithm
  File "/home/rm/wxz/EmbodiedAI/metaworld/metaworld-algorithms/metaworld_algorithms/rl/algorithms/base.py", line 8, in <module>
    import orbax.checkpoint as ocp
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/__init__.py", line 23, in <module>
    from orbax.checkpoint import aggregate_handlers
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/aggregate_handlers.py", line 24, in <module>
    from orbax.checkpoint._src.metadata import tree as tree_metadata
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/tree.py", line 36, in <module>
    from orbax.checkpoint._src.metadata import tree_rich_types
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/tree_rich_types.py", line 25, in <module>
    from orbax.checkpoint._src.metadata import value_metadata_entry
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/value_metadata_entry.py", line 25, in <module>
    from orbax.checkpoint._src.serialization import types
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/types.py", line 33, in <module>
    from orbax.checkpoint._src.serialization import serialization
  File "/home/rm/wxz/EmbodiedAI/metaworld/.meta_world_alg/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 403, in <module>
    dll: Optional[layout.DeviceLocalLayout],
                  ^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax.experimental.layout' has no attribute 'DeviceLocalLayout'
```

This seems like an issue with orbax-checkpoint. Which version of it are you using? Can you run the following command to find out: `pip freeze | grep orbax`? That would help me pin an earlier version in the repository's lockfile, thanks.

@rainx0r Thank you for your help. I ran the command ```pip freeze | grep orbax```, but there was no output. Then I tried``` pip install --upgrade jax jaxlib orbax-checkpoint ```, but it didn’t work either. Currently, when I run``` pip freeze | grep orbax```, my output is as follows:
```
orbax==0.1.9  
orbax-checkpoint==0.11.20  
```
And for ```pip freeze | grep jax:```
```
jax==0.6.2  
jaxlib==0.6.2  
```

Ok yeah looks like you're probably not using [uv](https://docs.astral.sh/uv/) or not installing the project with `pip install -e ".[cuda12]"` or similar, because in [pyproject.toml 's dependencies](https://github.com/rainx0r/metaworld-algorithms/blob/main/pyproject.toml#L14) there's `"orbax-checkpoint<=0.11.11"`.

 I highly suggest using `uv` if you can for many reasons. but in this case the project's stable dependencies are described in the auto generated `uv.lock` file so if you do `uv pip install -e ".[cuda12]"` ([as the README in this repo describes](https://github.com/rainx0r/metaworld-algorithms/tree/main?tab=readme-ov-file#from-a-clone-of-the-repository)) you should automatically get the right versions of all the dependencies.