Merge pull request #149 from jacobusmmsmit/fix-jtu-warning

jax -> jtu for relevant functions

Author: patrick-kidger

diffrax/adjoint.py

@@ -2,7 +2,6 @@
 from typing import Any, Dict
 
 import equinox as eqx
-import jax
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
@@ -171,7 +170,7 @@ def _loop_backsolve(y__args__terms, *, self, throw, init_state, **kwargs):
     del throw
     y, args, terms = y__args__terms
     init_state = eqx.tree_at(
-        lambda s: jax.tree_leaves(s.y), init_state, jax.tree_leaves(y)
+        lambda s: jtu.tree_leaves(s.y), init_state, jtu.tree_leaves(y)
     )
     del y
     return self._loop_fn(
@@ -409,7 +408,7 @@ def loop(self, *, args, terms, saveat, init_state, **kwargs):
         y = init_state.y
         sentinel = object()
         init_state = eqx.tree_at(
-            lambda s: jax.tree_leaves(s.y), init_state, replace_fn=lambda _: sentinel
+            lambda s: jtu.tree_leaves(s.y), init_state, replace_fn=lambda _: sentinel
         )
 
         final_state, aux_stats = _loop_backsolve(
@@ -421,7 +420,7 @@ def loop(self, *, args, terms, saveat, init_state, **kwargs):
         ys = final_state.ys
         final_state = jtu.tree_map(nondifferentiable_output, final_state)
         final_state = eqx.tree_at(
-            lambda s: jax.tree_leaves(s.ys), final_state, jax.tree_leaves(ys)
+            lambda s: jtu.tree_leaves(s.ys), final_state, jtu.tree_leaves(ys)
         )
 
         return final_state, aux_stats

diffrax/global_interpolation.py

@@ -547,7 +547,7 @@ def rectilinear_interpolation(
         out = jtu.tree_map(fn, replace_nans_at_start, ys)
     ys_treedef = jtu.tree_structure(ys)
     interp_treedef = jtu.tree_structure((0, 0))
-    return jax.tree_transpose(ys_treedef, interp_treedef, out)
+    return jtu.tree_transpose(ys_treedef, interp_treedef, out)
 
 
 def _hermite_forward(
@@ -725,4 +725,4 @@ def backward_hermite_coefficients(
             coeffs = jtu.tree_map(fn, ys, deriv0, replace_nans_at_start)
     ys_treedef = jtu.tree_structure(ys)
     coeffs_treedef = jtu.tree_structure((0, 0, 0, 0))
-    return jax.tree_transpose(ys_treedef, coeffs_treedef, coeffs)
+    return jtu.tree_transpose(ys_treedef, coeffs_treedef, coeffs)

diffrax/integrate.py

@@ -78,9 +78,9 @@ def _save(state: _State, t: Scalar) -> _State:
     save_index = save_index + 1
 
     return eqx.tree_at(
-        lambda s: [s.ts, s.save_index] + jax.tree_leaves(s.ys),
+        lambda s: [s.ts, s.save_index] + jtu.tree_leaves(s.ys),
         state,
-        [ts, save_index] + jax.tree_leaves(ys),
+        [ts, save_index] + jtu.tree_leaves(ys),
     )
 
 

diffrax/solver/milstein.py

@@ -183,7 +183,7 @@ def step(
                     leaf = leaf - Δt * eye
                 leaves_ΔwΔw.append(leaf)
         tree_ΔwΔw = tree_Δw.compose(tree_Δw)
-        ΔwΔw = jax.tree_unflatten(tree_ΔwΔw, leaves_ΔwΔw)
+        ΔwΔw = jtu.tree_unflatten(tree_ΔwΔw, leaves_ΔwΔw)
         # ΔwΔw has structure (tree(Δw), tree(Δw), leaf(Δw), leaf(Δw))
 
         #
@@ -250,11 +250,11 @@ def _to_treemap(_Δw, _g0):
         # Which we now transform into its isomorphic matrix form, as above.
         g0_matrix = jax.jacfwd(lambda _Δw: diffusion.prod(g0, _Δw))(Δw)
         # g0_matrix has structure (tree(y0), tree(Δw), leaf(y0), leaf(Δw))
-        g0_matrix = jax.tree_transpose(y_treedef, Δw_treedef, g0_matrix)
+        g0_matrix = jtu.tree_transpose(y_treedef, Δw_treedef, g0_matrix)
         # g0_matrix has structure (tree(Δw), tree(y0), leaf(y0), leaf(Δw))
         v0_matrix = jtu.tree_map(_to_treemap, Δw, g0_matrix)
         # v0_matrix has structure (tree(Δw), tree(y0), tree(Δw), leaf(y0), leaf(Δw), leaf(Δw))  # noqa: E501
-        v0_matrix = jax.tree_transpose(
+        v0_matrix = jtu.tree_transpose(
             Δw_treedef, y_treedef.compose(Δw_treedef), v0_matrix
         )
         # v0_matrix has structure (tree(y0), tree(Δw), tree(Δw), leaf(y0), leaf(Δw), leaf(Δw))  # noqa: E501
@@ -275,7 +275,7 @@ def _dot(_, _v0):
             # ΔwΔw has structure (tree(Δw), tree(Δw), leaf(Δw), leaf(Δw))
             _dotted = jtu.tree_map(__dot, _v0, ΔwΔw)
             # _dotted has structure (tree(Δw), tree(Δw), leaf(y0))
-            _out = sum(jax.tree_leaves(_dotted))
+            _out = sum(jtu.tree_leaves(_dotted))
             # _out has structure (leaf(y0),)
             return _out
 

diffrax/term.py

@@ -383,7 +383,7 @@ def is_vf_expensive(
         args: PyTree,
     ) -> bool:
         control = self.contr(t0, t1)
-        if sum(c.size for c in jax.tree_leaves(control)) in (0, 1):
+        if sum(c.size for c in jtu.tree_leaves(control)) in (0, 1):
             return False
         else:
             return True
@@ -412,8 +412,8 @@ def vf(
         # PyTree structure. (This is because `self.vf_prod` is linear in `control`.)
         control = self.contr(t, t)
 
-        y_size = sum(np.size(yi) for yi in jax.tree_leaves(y))
-        control_size = sum(np.size(ci) for ci in jax.tree_leaves(control))
+        y_size = sum(np.size(yi) for yi in jtu.tree_leaves(y))
+        control_size = sum(np.size(ci) for ci in jtu.tree_leaves(control))
         if y_size > control_size:
             make_jac = jax.jacfwd
         else:
@@ -441,7 +441,7 @@ def _fn(_control):
             raise NotImplementedError(
                 "`AdjointTerm.vf` not implemented for `None` controls or states."
             )
-        return jax.tree_transpose(vf_prod_tree, control_tree, jac)
+        return jtu.tree_transpose(vf_prod_tree, control_tree, jac)
 
     def contr(self, t0: Scalar, t1: Scalar) -> PyTree:
         return self.term.contr(t0, t1)
@@ -467,16 +467,16 @@ def _get_vf_tree(_, tree):
         jtu.tree_map(_get_vf_tree, control, vf)
         assert vf_prod_tree is not sentinel
 
-        vf = jax.tree_transpose(control_tree, vf_prod_tree, vf)
+        vf = jtu.tree_transpose(control_tree, vf_prod_tree, vf)
 
-        example_vf_prod = jax.tree_unflatten(
+        example_vf_prod = jtu.tree_unflatten(
             vf_prod_tree, [0 for _ in range(vf_prod_tree.num_leaves)]
         )
 
         def _contract(_, vf_piece):
             assert jtu.tree_structure(vf_piece) == control_tree
             _contracted = jtu.tree_map(_prod, vf_piece, control)
-            return sum(jax.tree_leaves(_contracted), 0)
+            return sum(jtu.tree_leaves(_contracted), 0)
 
         return jtu.tree_map(_contract, example_vf_prod, vf)
 

test/helpers.py

@@ -5,7 +5,6 @@
 
 import diffrax
 import equinox as eqx
-import jax
 import jax.numpy as jnp
 import jax.random as jrandom
 import jax.tree_util as jtu
@@ -59,7 +58,7 @@ def random_pytree(key, treedef):
         dim_sizes = jrandom.randint(sizekey, (num_dims,), 0, 5)
         value = jrandom.normal(valuekey, dim_sizes)
         leaves.append(value)
-    return jax.tree_unflatten(treedef, leaves)
+    return jtu.tree_unflatten(treedef, leaves)
 
 
 treedefs = [