Replace deprecated jax.tree_map function with jax.tree_util.tree_map

sotetsuk / pgx

♟️ Vectorized RL game environments in JAX

http://sotets.uk/pgx/

Apache License 2.0

372 stars 23 forks source link

Closed KazukiOhta closed 1 month ago

KazukiOhta commented 1 month ago

The use of jax.tree_map triggers the following warning.

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).

Therefore jax.tree_map is replaced by jax.tree_util.tree_map which is available in any JAX version.

sotetsuk commented 1 month ago

Thank you for your contribution! 🙏

KazukiOhta commented 1 month ago

Thank you too for providing a nice implementation of Gumbel AlphaZero.