New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

simple_cnn baseline doesn't work with jax nightly #1071

Open

sycamoreoak opened this issue May 29, 2024 · 0 comments

sycamoreoak commented May 29, 2024

python3 ./simple_cnn.py 
Traceback (most recent call last):
  File "~/.local/lib/python3.10/site-packages/flax/core/axes_scan.py", line 21, in <module>
    from jax import linear_util as lu
ImportError: cannot import name 'linear_util' from 'jax' (~/.local/lib/python3.10/site-packages/jax/__init__.py)

If I modify the linear_util line in to this:

from jax.extend import linear_util as lu

then I get:

  File "~/.local/lib/python3.10/site-packages/flax/struct.py", line 141, in dataclass
    if tuple(map(int, jax.version.__version__.split('.'))) >= (0, 3, 1):
ValueError: invalid literal for int() with base 10: 'dev20240528'

The text was updated successfully, but these errors were encountered:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment