AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'

[zhongJY1998]

Hello, please tell me about the error when running vision_chat.py, where jax==0.4.23, tux==0.0.2 The specific errors are as follows: File "/home/LWM-main/lwm/vision_chat.py", line 12, in from tux import ( File "/root/miniconda3/envs/lwm/lib/python3.10/site-packages/tux/init.py", line 17, in from .optimizers import (AdamWOptimizerFactory, get_weight_decay_mask, optax_add_scheduled_weight_decay, File "/root/miniconda3/envs/lwm/lib/python3.10/site-packages/tux/optimizers.py", line 193, in class OptaxScheduledWeightDecayState(NamedTuple): File "/root/miniconda3/envs/lwm/lib/python3.10/site-packages/tux/optimizers.py", line 194, in OptaxScheduledWeightDecayState count: jnp.DeviceArray File "/root/miniconda3/envs/lwm/lib/python3.10/site-packages/jax/_src/deprecations.py", line 53, in getattr raise AttributeError(f"module {module!r} has no attribute {name!r}") AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'

[ersanliqiao]

same error