Use `jax.tree.map` because `jax.tree.map` is deprecated

RobertTLange / gymnax

RL Environments in JAX 🌍

Apache License 2.0

577 stars 54 forks source link

Use `jax.tree.map` because `jax.tree.map` is deprecated #70

Open cool-RR opened 2 months ago

cool-RR commented 2 months ago

<ipython-input-22-bfd27f6f2527>:8: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).