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: 634281277
Change-Id: I9a440f4625a239ede5982116a059a496a9a6b150

ferminet/utils/multi_host.py

@@ -30,8 +30,8 @@ def check_synced(obj, name):
     True if object is in sync across all devices and False otherwise.
   """
   for i in range(1, jax.local_device_count()):
-    norms = jax.tree_map(lambda x: jnp.linalg.norm(x[0] - x[i]), obj)  # pylint: disable=cell-var-from-loop
-    total_norms = sum(jax.tree_leaves(norms))
+    norms = jax.tree.map(lambda x: jnp.linalg.norm(x[0] - x[i]), obj)  # pylint: disable=cell-var-from-loop
+    total_norms = sum(jax.tree.leaves(norms))
     if total_norms != 0.0:
       logging.info(
           '%s object is not synced across device 0 and %d. The total norm'