Merge pull request #19958 from jakevdp:jax-tree

PiperOrigin-RevId: 610533158

Author: jax authors

tests/api_test.py

@@ -1042,7 +1042,7 @@ def f(x):
       self.assertEqual(
           obj.in_avals,
           ((core.ShapedArray([], expected_dtype, weak_type=True),), {}))
-      self.assertEqual(obj.in_tree, jax.tree_util.tree_flatten(((0,), {}))[1])
+      self.assertEqual(obj.in_tree, jax.tree.flatten(((0,), {}))[1])
 
   def test_jit_lower_duck_typing(self):
     f_jit = jit(lambda x: 2 * x)
@@ -2490,7 +2490,7 @@ def fun(x, y):
     x = (jnp.ones(2), jnp.ones(2))
     y = 3.
     out_shape = api.eval_shape(fun, x, y)
-    out_shape = tree_util.tree_map(np.shape, out_shape)
+    out_shape = jax.tree.map(np.shape, out_shape)
 
     self.assertEqual(out_shape, {'hi': (2,)})
 
@@ -3004,7 +3004,7 @@ def test_vmap_in_axes_tree_prefix_error(self):
         ValueError,
         "vmap in_axes specification must be a tree prefix of the corresponding "
         r"value, got specification \(\[0\],\) for value tree "
-        + re.escape(f"{tree_util.tree_structure((value_tree,))}."),
+        + re.escape(f"{jax.tree.structure((value_tree,))}."),
         lambda: api.vmap(lambda x: x, in_axes=([0],))(value_tree)
     )
 
@@ -7013,8 +7013,8 @@ def foo_jvp(primals, tangents):
             "must produce primal and tangent outputs "
             "with equal container (pytree) structures, but got "
             "{} and {} respectively.".format(
-                tree_util.tree_structure((1,)),
-                tree_util.tree_structure([1, 2]))
+                jax.tree.structure((1,)),
+                jax.tree.structure([1, 2]))
         ),
         lambda: api.jvp(f, (2.,), (1.,)))
 
@@ -7729,9 +7729,9 @@ def _unpack(x):
 
     def _vmap(fun):
       def _fun(*args):
-        args = tree_util.tree_map(_pack, args)
+        args = jax.tree.map(_pack, args)
         out = jax.vmap(fun)(*args)
-        out = tree_util.tree_map(_unpack, out)
+        out = jax.tree.map(_unpack, out)
         return out
       return _fun
 
@@ -8242,8 +8242,8 @@ def foo_bwd(_, g):
             "and in particular must produce a tuple of length equal to the "
             "number of arguments to the primal function, but got VJP output "
             "structure {} for primal input structure {}.".format(
-                tree_util.tree_structure((1, 1)),
-                tree_util.tree_structure((1,)))
+                jax.tree.structure((1, 1)),
+                jax.tree.structure((1,)))
         ),
         lambda: api.grad(f)(2.))
 
@@ -9017,9 +9017,9 @@ def _unpack(x):
 
     def _vmap(fun):
       def _fun(*args):
-        args = tree_util.tree_map(_pack, args)
+        args = jax.tree.map(_pack, args)
         out = jax.vmap(fun)(*args)
-        out = tree_util.tree_map(_unpack, out)
+        out = jax.tree.map(_unpack, out)
         return out
       return _fun
 
@@ -9281,7 +9281,7 @@ def custom_transpose(example_out):
     return _custom_transpose(out_type, example_out)
   return partial(
       _custom_transpose,
-      tree_util.tree_map(
+      jax.tree.map(
           lambda x: core.get_aval(x).at_least_vspace(), example_out))
 
 
@@ -10139,21 +10139,21 @@ def rule(axis_size, in_batched, xs):
             f_jvp, in_axes=(None, 0), out_axes=(None, 0))(x, txs))
 
   def test_tree(self):
-    tree_sin = partial(tree_util.tree_map, jnp.sin)
-    tree_cos = partial(tree_util.tree_map, jnp.cos)
+    tree_sin = partial(jax.tree.map, jnp.sin)
+    tree_cos = partial(jax.tree.map, jnp.cos)
 
     x, xs = jnp.array(1.), jnp.arange(3)
     x  = (x,  [x  + 1, x  + 2], [x  + 3], x  + 4)
     xs = (xs, [xs + 1, xs + 2], [xs + 3], xs + 4)
-    in_batched_ref = tree_util.tree_map(lambda _: True, x)
+    in_batched_ref = jax.tree.map(lambda _: True, x)
 
     @jax.custom_batching.custom_vmap
     def f(xs): return tree_sin(xs)
 
     @f.def_vmap
     def rule(axis_size, in_batched, xs):
       self.assertEqual(in_batched, [in_batched_ref])
-      sz, = {z.shape[0] for z in tree_util.tree_leaves(xs)}
+      sz, = {z.shape[0] for z in jax.tree.leaves(xs)}
       self.assertEqual(axis_size, sz)
       return tree_cos(xs), in_batched[0]
 
@@ -10163,21 +10163,21 @@ def rule(axis_size, in_batched, xs):
     self.assertAllClose(ys, tree_cos(xs))
 
   def test_tree_with_nones(self):
-    tree_sin = partial(tree_util.tree_map, jnp.sin)
-    tree_cos = partial(tree_util.tree_map, jnp.cos)
+    tree_sin = partial(jax.tree.map, jnp.sin)
+    tree_cos = partial(jax.tree.map, jnp.cos)
 
     x, xs = jnp.array(1.), jnp.arange(3)
     x  = (x,  [x  + 1, None], [x  + 3], None)
     xs = (xs, [xs + 1, None], [xs + 3], None)
-    in_batched_ref = tree_util.tree_map(lambda _: True, x)
+    in_batched_ref = jax.tree.map(lambda _: True, x)
 
     @jax.custom_batching.custom_vmap
     def f(xs): return tree_sin(xs)
 
     @f.def_vmap
     def rule(axis_size, in_batched, xs):
       self.assertEqual(in_batched, [in_batched_ref])
-      sz, = {z.shape[0] for z in tree_util.tree_leaves(xs)}
+      sz, = {z.shape[0] for z in jax.tree.leaves(xs)}
       self.assertEqual(axis_size, sz)
       return tree_cos(xs), in_batched[0]
 

tests/core_test.py

@@ -28,7 +28,6 @@
 from jax import jvp, linearize, vjp, jit, make_jaxpr
 from jax.api_util import flatten_fun_nokwargs
 from jax import config
-from jax.tree_util import tree_flatten, tree_unflatten, tree_map, tree_reduce
 
 from jax._src import core
 from jax._src import linear_util as lu
@@ -49,17 +48,17 @@ def call(f, *args):
 
 @util.curry
 def core_call(f, *args):
-  args, in_tree = tree_flatten(args)
+  args, in_tree = jax.tree.flatten(args)
   f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
   out = core.call_p.bind(f, *args)
-  return tree_unflatten(out_tree(), out)
+  return jax.tree.unflatten(out_tree(), out)
 
 @util.curry
 def core_closed_call(f, *args):
-  args, in_tree = tree_flatten(args)
+  args, in_tree = jax.tree.flatten(args)
   f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
   out = core.closed_call_p.bind(f, *args)
-  return tree_unflatten(out_tree(), out)
+  return jax.tree.unflatten(out_tree(), out)
 
 def simple_fun(x, y):
   return jnp.sin(x * y)
@@ -175,24 +174,24 @@ def test_tree_map(self):
     zs = ({'a': 11}, [22, 33])
 
     f = lambda x, y: x + y
-    assert tree_map(f, xs, ys) == zs
+    assert jax.tree.map(f, xs, ys) == zs
     try:
-      tree_map(f, xs, ys_bad)
+      jax.tree.map(f, xs, ys_bad)
       assert False
     except (TypeError, ValueError):
       pass
 
   def test_tree_flatten(self):
-    flat, _ = tree_flatten(({'a': 1}, [2, 3], 4))
+    flat, _ = jax.tree.flatten(({'a': 1}, [2, 3], 4))
     assert flat == [1, 2, 3, 4]
 
   def test_tree_unflatten(self):
     tree = [(1, 2), {"roy": (3, [4, 5, ()])}]
-    flat, treedef = tree_flatten(tree)
+    flat, treedef = jax.tree.flatten(tree)
     assert flat == [1, 2, 3, 4, 5]
-    tree2 = tree_unflatten(treedef, flat)
-    nodes_equal = tree_map(operator.eq, tree, tree2)
-    assert tree_reduce(operator.and_, nodes_equal)
+    tree2 = jax.tree.unflatten(treedef, flat)
+    nodes_equal = jax.tree.map(operator.eq, tree, tree2)
+    assert jax.tree.reduce(operator.and_, nodes_equal)
 
   @jtu.sample_product(
       dtype=[*jtu.dtypes.all, object, [('i', 'i4'), ('f', 'f4')]]

tests/custom_linear_solve_test.py

@@ -25,7 +25,6 @@
 from jax import lax
 from jax.ad_checkpoint import checkpoint
 from jax._src import test_util as jtu
-from jax import tree_util
 import jax.numpy as jnp  # scan tests use numpy
 import jax.scipy as jsp
 
@@ -160,7 +159,7 @@ def linear_solve_aux(a, b):
     # vmap test
     c = rng.randn(3, 2)
     expected = jnp.linalg.solve(a, c)
-    expected_aux = tree_util.tree_map(partial(np.repeat, repeats=2), array_aux)
+    expected_aux = jax.tree.map(partial(np.repeat, repeats=2), array_aux)
     actual_vmap, vmap_aux = jax.vmap(linear_solve_aux, (None, 1), -1)(a, c)
 
     self.assertAllClose(expected, actual_vmap)
@@ -473,7 +472,7 @@ def solve(mv, b):
       return mv(b), aux
 
     def solve_aux(x):
-      matvec = lambda y: tree_util.tree_map(partial(jnp.dot, A), y)
+      matvec = lambda y: jax.tree.map(partial(jnp.dot, A), y)
       return lax.custom_linear_solve(matvec, (x, x), solve, solve, symmetric=True, has_aux=True)
 
     rng = self.rng()

tests/custom_root_test.py

@@ -22,7 +22,6 @@
 import jax
 from jax import lax
 from jax._src import test_util as jtu
-from jax import tree_util
 import jax.numpy as jnp  # scan tests use numpy
 import jax.scipy as jsp
 
@@ -227,7 +226,7 @@ def pos_def_solve(g, b):
     expected_fwd_val = expected_fwd(a, b)
     self.assertAllClose(fwd_val, expected_fwd_val, rtol={np.float32: 5E-6, np.float64: 5E-12})
 
-    jtu.check_close(fwd_aux, tree_util.tree_map(jnp.zeros_like, fwd_aux))
+    jtu.check_close(fwd_aux, jax.tree.map(jnp.zeros_like, fwd_aux))
 
   def test_custom_root_errors(self):
     with self.assertRaisesRegex(TypeError, re.escape("f() output pytree")):

tests/export_harnesses_multi_platform_test.py

@@ -155,7 +155,7 @@ def export_and_compare_to_native(
       if device.platform in skip_run_on_platforms:
         logging.info("Skipping running on %s", device)
         continue
-      device_args = jax.tree_util.tree_map(
+      device_args = jax.tree.map(
           lambda x: jax.device_put(x, device), args
       )
       logging.info("Running harness natively on %s", device)

tests/export_test.py

@@ -25,7 +25,6 @@
 import jax
 from jax import lax
 from jax import numpy as jnp
-from jax import tree_util
 from jax.experimental import export
 from jax.experimental.export import _export
 from jax.experimental import pjit
@@ -185,7 +184,7 @@ def my_fun(x):
     self.assertEqual("my_fun", exp.fun_name)
     self.assertEqual((export.default_lowering_platform(),),
                      exp.lowering_platforms)
-    self.assertEqual(tree_util.tree_flatten(((1,), {}))[1], exp.in_tree)
+    self.assertEqual(jax.tree.flatten(((1,), {}))[1], exp.in_tree)
     self.assertEqual((core.ShapedArray((4,), dtype=np.float32),), exp.in_avals)
     self.assertEqual((core.ShapedArray((4,), dtype=np.float32),), exp.out_avals)
 
@@ -201,9 +200,9 @@ def f(a_b_pair, *, a, b):
     self.assertEqual(exp.lowering_platforms, ("cpu",))
     args = ((a, b),)
     kwargs = dict(a=a, b=b)
-    self.assertEqual(exp.in_tree, tree_util.tree_flatten((args, kwargs))[1])
+    self.assertEqual(exp.in_tree, jax.tree.flatten((args, kwargs))[1])
     self.assertEqual(exp.in_avals, (a_aval, b_aval, a_aval, b_aval))
-    self.assertEqual(exp.out_tree, tree_util.tree_flatten(f(*args, **kwargs))[1])
+    self.assertEqual(exp.out_tree, jax.tree.flatten(f(*args, **kwargs))[1])
     self.assertEqual(exp.out_avals, (a_aval, b_aval, a_aval, b_aval, a_aval, b_aval))
 
   def test_basic(self):

tests/filecheck/jax_filecheck_helpers.py

@@ -15,16 +15,15 @@
 # Helpers for writing JAX filecheck tests.
 
 import jax
-import jax.tree_util as tree_util
 import numpy as np
 
 def print_ir(*prototypes):
   def lower(f):
     """Prints the MLIR IR that results from lowering `f`.
 
     The arguments to `f` are taken to be arrays shaped like `prototypes`."""
-    inputs = tree_util.tree_map(np.array, prototypes)
-    flat_inputs, _ = tree_util.tree_flatten(inputs)
+    inputs = jax.tree.map(np.array, prototypes)
+    flat_inputs, _ = jax.tree.flatten(inputs)
     shape_strs = " ".join([f"{x.dtype.name}[{','.join(map(str, x.shape))}]"
                            for x in flat_inputs])
     name = f.func.__name__ if hasattr(f, "func") else f.__name__

tests/host_callback_test.py

@@ -34,7 +34,6 @@
 from jax import dtypes
 from jax import lax
 from jax import numpy as jnp
-from jax import tree_util
 from jax.experimental import host_callback as hcb
 from jax.experimental import pjit
 from jax.sharding import PartitionSpec as P
@@ -887,7 +886,7 @@ def func(x):
     # making the Jaxpr does not print anything
     hcb.barrier_wait()
 
-    treedef = tree_util.tree_structure(arg)
+    treedef = jax.tree.structure(arg)
     assertMultiLineStrippedEqual(
         self, f"""
       {{ lambda ; a:f32[]. let
@@ -1027,7 +1026,7 @@ def make_ct(res):
           return res
         ct_dtype = core.primal_dtype_to_tangent_dtype(res_dtype)
         return np.ones(np.shape(res), dtype=ct_dtype)
-      cts = tree_util.tree_map(make_ct, res_f_of_args)
+      cts = jax.tree.map(make_ct, res_f_of_args)
       def f_vjp(args, cts):
         res, pullback = jax.vjp(f, *args)
         return pullback(cts)

tests/infeed_test.py

@@ -72,7 +72,7 @@ def f(x):
 
     device = jax.local_devices()[0]
     # We must transfer the flattened data, as a tuple!!!
-    flat_to_infeed, _ = jax.tree_util.tree_flatten(to_infeed)
+    flat_to_infeed, _ = jax.tree.flatten(to_infeed)
     device.transfer_to_infeed(tuple(flat_to_infeed))
     self.assertAllClose(f(x), to_infeed)
 

tests/jet_test.py

@@ -58,8 +58,8 @@ class JetTest(jtu.JaxTestCase):
   def check_jet(self, fun, primals, series, atol=1e-5, rtol=1e-5,
                 check_dtypes=True):
     # Convert to jax arrays to ensure dtype canonicalization.
-    primals = jax.tree_util.tree_map(jnp.asarray, primals)
-    series = jax.tree_util.tree_map(jnp.asarray, series)
+    primals = jax.tree.map(jnp.asarray, primals)
+    series = jax.tree.map(jnp.asarray, series)
 
     y, terms = jet(fun, primals, series)
     expected_y, expected_terms = jvp_taylor(fun, primals, series)
@@ -73,8 +73,8 @@ def check_jet(self, fun, primals, series, atol=1e-5, rtol=1e-5,
   def check_jet_finite(self, fun, primals, series, atol=1e-5, rtol=1e-5,
                        check_dtypes=True):
     # Convert to jax arrays to ensure dtype canonicalization.
-    primals = jax.tree_util.tree_map(jnp.asarray, primals)
-    series = jax.tree_util.tree_map(jnp.asarray, series)
+    primals = jax.tree.map(jnp.asarray, primals)
+    series = jax.tree.map(jnp.asarray, series)
 
     y, terms = jet(fun, primals, series)
     expected_y, expected_terms = jvp_taylor(fun, primals, series)