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: 634381464

notebooks/jax_raycast.ipynb

@@ -815,7 +815,7 @@
         "\n",
         "def render_frame(t):\n",
         "  t = t*t*(3-2*t)  # easing\n",
-        "  s = jax.tree_map(partial(cubic, t), s0, s1, s2, s0)\n",
+        "  s = jax.tree.map(partial(cubic, t), s0, s1, s2, s0)\n",
         "  return render_scene(**s)\n",
         "\n",
         "animate(render_frame, 10)"
@@ -857,7 +857,7 @@
     {
       "cell_type": "markdown",
       "source": [
-        "The code above uses `jax.tree_map` to conveniently interpolate between nested data structures that contain information about the scene, including light and camera position.\n",
+        "The code above uses `jax.tree.map` to conveniently interpolate between nested data structures that contain information about the scene, including light and camera position.\n",
         "\n",
         "Rendering a 10 second 60 fps animation takes from ~20 to less than 10 seconds depending on the available GPU (I tried V100 and A100, also running the first time in a Colab session takes a bit longer). Not bad for a ~100 lines of code engine built from scratch on top of the generic array processing language.\n",
         "\n",