Replace deprecated jax.tree_* functions with jax.tree.*

The top-level `jax.tree_*` aliases have long been deprecated, and will soon be removed. Alternate APIs are in `jax.tree_util`, with shorter aliases in the `jax.tree` submodule, added in JAX version 0.4.25. PiperOrigin-RevId: 634020884
google-research · May 15, 2024 · e588c15 · e588c15
1 parent d392e41
commit e588c15
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/diffren/jax/utils/mesh.py b/diffren/jax/utils/mesh.py
@@ -129,7 +129,7 @@ def remove_unused_vertices(vertices, attributes, triangles):
       .set(True)
   )
   kept_vertices = vertices[vertex_mask, :]
-  kept_attributes = jax.tree_map(lambda x: x[vertex_mask, :], attributes)
+  kept_attributes = jax.tree.map(lambda x: x[vertex_mask, :], attributes)
 
   vertex_mask_offset = jnp.pad(vertex_mask, ((1, 0)), constant_values=False)
   new_vert_indices = jnp.cumsum(vertex_mask_offset)