Error importing ciclo

[hylkedonker]

Hi,

Your package looks pretty neat, thanks for sharing!
However, I haven't managed to get it to work. When I run

import ciclo

I get the error message

Traceback (most recent call last):
  File "/home/user/myfile.py", line 3, in <module>
    import ciclo
  File "/home/user/.local/lib/python3.11/site-packages/ciclo/__init__.py", line 5, in <module>
    from .callbacks import checkpoint as checkpoint
  File "/home/user/.local/lib/python3.11/site-packages/ciclo/callbacks.py", line 167, in <module>
    from flax.training import checkpoints as flax_checkpoints
  File "/home/user/.local/lib/python3.11/site-packages/flax/training/checkpoints.py", line 38, in <module>
    from flax.training import orbax_utils
  File "/home/user/.local/lib/python3.11/site-packages/flax/training/orbax_utils.py", line 25, in <module>
    import orbax.checkpoint as ocp
  File "/home/user/.local/lib/python3.11/site-packages/orbax/checkpoint/__init__.py", line 20, in <module>
    from orbax.checkpoint import checkpoint_utils
  File "/home/user/.local/lib/python3.11/site-packages/orbax/checkpoint/checkpoint_utils.py", line 25, in <module>
    from orbax.checkpoint import type_handlers
  File "/home/user/.local/lib/python3.11/site-packages/orbax/checkpoint/type_handlers.py", line 24, in <module>
    from jax.experimental.gda_serialization import serialization
ModuleNotFoundError: No module named 'jax.experimental.gda_serialization

I have the following package versions installed:

flax==0.7.4
jax==0.4.19
jax-metrics==0.2.5
jaxlib==0.4.19
jaxtyping==0.2.23
orbax==0.1.3
orbax-checkpoint==0.4.1

This looks like it is related to the outdated orbax version (0.1.3 instead of 0.1.9) that clashes with flax. When I update orbax (pip3 install --upgrade orbax) I get a slightly different error:

2023-10-27 14:00:12.914634: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
  File "/home/user/model.py", line 3, in <module>
    import ciclo
  File "/home/user/.local/lib/python3.11/site-packages/ciclo/__init__.py", line 5, in <module>
    from .callbacks import checkpoint as checkpoint
  File "/home/user/.local/lib/python3.11/site-packages/ciclo/callbacks.py", line 167, in <module>
    from flax.training import checkpoints as flax_checkpoints
  File "/home/user/.local/lib/python3.11/site-packages/flax/training/checkpoints.py", line 589, in <module>
    orbax_checkpointer: Optional[ocp.Checkpointer] = None,
                                 ^^^^^^^^^^^^^^^^
AttributeError: module 'orbax.checkpoint' has no attribute 'Checkpointer'

Any thoughts how I can get ciclo to work without downgrading to very old versions of jax and flax?

Thanks in advance,

Hylke

[cgarciae]

Hey @hylkedonker ! Can you try upgrading to the new version of Flax?

[hylkedonker]

Yes, same problem on flax==0.7.5.
Strangely enough, I haven't been able to replicate my problem on Ubuntu 23.10 Docker. That is, the following seems to work fine, although it should mirror (atleast, closely) my system:

docker run ubuntu:23.10 sh -c 'apt update \
&& apt install --assume-yes python3-pip git \
&& pip3 install --break-system-packages flax \
&& pip3 install --break-system-packages git+https://github.com/cgarciae/ciclo.git \
&& pip3 install --upgrade --break-system-packages orbax \
&& python3 -c "import ciclo"'

Any thoughts how to dig further?

[cgarciae]

Hey, does the pypi package work?

[hylkedonker]

Nope, same problem as HEAD.
Anything that I can do to further clarify?