英文:
No module named 'jax.experimental.global_device_array' when running the official Flax Example on Colab with V100
问题
以下是您要翻译的内容：
"我一直在尝试理解这个官方的flax示例，基于一个带有V100的Coalb pro+账户。当我执行命令 python main.py --workdir=./imagenet --config=configs/v100_x8.py 时，返回的错误是：
File "/content/FlaxImageNet/main.py", line 29, in <module>
import train
File "/content/FlaxImageNet/train.py", line 30, in <module>
from flax.training import checkpoints
File "/usr/local/lib/python3.10/dist-packages/flax/training/checkpoints.py", line 34, in <module>
from jax.experimental.global_device_array import GlobalDeviceArray
ModuleNotFoundError: No module named 'jax.experimental.global_device_array'
我不确定 global_device_array 是否已从 jax.experimental 包中移动，或者它是否不再需要，或者是否已被其他等效方法替代。
英文:
I have been trying to understand this official flax example, based on a Coalb pro+ account with V100. When I execute the command python main.py --workdir=./imagenet --config=configs/v100_x8.py, the returned error is
File "/content/FlaxImageNet/main.py", line 29, in <module>
import train
File "/content/FlaxImageNet/train.py", line 30, in <module>
from flax.training import checkpoints
File "/usr/local/lib/python3.10/dist-packages/flax/training/checkpoints.py", line 34,
in <module>
from jax.experimental.global_device_array import GlobalDeviceArray
ModuleNotFoundError: No module named 'jax.experimental.global_device_array'
I am not sure whether global_device_array has been moved from jax.experimental package or it is no longer needed or replaced by other equivalent methods.
答案1
得分: 1
GlobalDeviceArray 在 JAX 版本 0.4.1 中被弃用，并在 JAX 版本 0.4.7 中移除。
考虑到这一点，似乎相关代码需要 JAX 版本 0.4.6 或更早的版本。您可以考虑向 flax 项目报告这个不兼容性：http://github.com/google/flax/。
英文:
GlobalDeviceArray was deprecated in JAX version 0.4.1 and removed in JAX version 0.4.7.
With that in mind, it seems the code in question requires JAX version 0.4.6 or older. You might consider reporting this incompatibility to the flax project: http://github.com/google/flax/.
通过集体智慧和协作来改善编程学习和解决问题的方式。致力于成为全球开发者共同参与的知识库，让每个人都能够通过互相帮助和分享经验来进步。
评论