jax.interpreters.xla.DeviceArray deprecated in jax 0.4.11

[tisabe]

Hi all,
I noticed after updating jax to the newest version (0.4.11) I get an error during import of optax:

AttributeError: module 'jax.interpreters.xla' has no attribute 'DeviceArray'

After looking into the changelog of jax, it seems jax.interpreters.xla.DeviceArray is now deprecated and jax.Array should be used.

Edit: Seems hard to reproduce, I installed optax in a new virtual environment where it works fine. Still not working in the old venv.

[hbq1]

Hi @tisabe, it must have been resolved in the latest versions of the optax's dependencies.
Let me close this issue for now, but feel free to re-open if you encounter this error again.

[davzuliani]

Hi,
I'm sorry for reopening this issue, but I just encountered it.
What is the solution?

Thanks!

[tisabe]

Hi @davzuliani
For me creating a fresh venv and installing the newest versions of jax and optax worked.

[davzuliani]

Hi @tisabe, thanks for the tip.
I guess I should go with venv, I just tried with conda env and it's still failing

[tisabe]

Okay, I tracked down where the issue actually came from. Looks like the chex requirement should be updated to >=v0.1.6 or >=v0.1.7. After updating chex, this error should be fixed. @hbq1 should this be opened in a new issue?

[louieworth]

Hi, I also meet the same error when I am trying to save my model. Do you know how to solve it?

def save(self, save_path: str):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    with open(save_path, 'wb') as f:
        f.write(flax.serialization.to_bytes(self.params))

JAX version: 0.4.12
JAXlib: 0.4.12

error track:

module 'jax.interpreters.xla' has no attribute 'DeviceArray'
  File "common.py", line 95, in save
    f.write(flax.serialization.to_bytes(self.params))
AttributeError: module 'jax.interpreters.xla' has no attribute 'DeviceArray'

[bainro]

pip install chex==0.1.7 --force-reinstall worked for me!