ci: fix tests in pipeline

instadeepai / jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

https://instadeepai.github.io/jumanji

Apache License 2.0

607 stars 77 forks source link

Closed clement-bonnet closed 7 months ago

clement-bonnet commented 7 months ago

jax.random.KeyArray -> jax.random.PRNGKey
use jnp.isclose instead of == for assertions with float32
convert MMST's frac variable to int32 because numpy uses float64 by default and Jax does not.