instadeepai / jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX
https://instadeepai.github.io/jumanji
Apache License 2.0
584 stars 71 forks source link

fix: jax version constraint #159

Closed aar65537 closed 1 year ago

aar65537 commented 1 year ago

Description

The max jax version is set to 0.4.1 because of the deprecation of DeviceArray (see pr #150). However, DeviceArray was deprecated in version 0.4.11, so the max jax version should be 0.4.10.

What Jumanji version are you using?

jumanji 0.2.2

Which accelerator(s) are you using?

GPU

Additional System Info

3.8

Additional Context

No response

(Optional) Suggestion

Change

jax>=0.2.26,<=0.4.1
jaxlib>=0.1.74,<=0.4.1

to

jax>=0.2.26,<=0.4.10
jaxlib>=0.1.74,<=0.4.10

in requirements/requirements.txt