Updated to latest Equinox (v0.10.8)

Author: patrick-kidger

README.md

@@ -21,7 +21,7 @@ _From a technical point of view, the internal structure of the library is pretty
 pip install diffrax
 ```
 
-Requires Python 3.9+, JAX 0.4.4+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.10.4+.
+Requires Python 3.9+, JAX 0.4.13+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.10.8+.
 
 ## Documentation
 

diffrax/adjoint.py

@@ -493,7 +493,9 @@ def _loop_backsolve(y__args__terms, *, self, throw, init_state, **kwargs):
     )
 
 
-def _loop_backsolve_fwd(y__args__terms, **kwargs):
+@_loop_backsolve.def_fwd
+def _loop_backsolve_fwd(perturbed, y__args__terms, **kwargs):
+    del perturbed
     final_state, aux_stats = _loop_backsolve(y__args__terms, **kwargs)
     # Note that `final_state.save_state` has type `PyTree[SaveState]`; here we are
     # relying on the guard in `BacksolveAdjoint` that it have trivial structure.
@@ -502,9 +504,18 @@ def _loop_backsolve_fwd(y__args__terms, **kwargs):
     return (final_state, aux_stats), (ts, ys)
 
 
+def _materialise_none(y, grad_y):
+    if grad_y is None and eqx.is_inexact_array(y):
+        return jnp.zeros_like(y)
+    else:
+        return grad_y
+
+
+@_loop_backsolve.def_bwd
 def _loop_backsolve_bwd(
     residuals,
     grad_final_state__aux_stats,
+    perturbed,
     y__args__terms,
     *,
     self,
@@ -525,13 +536,15 @@ def _loop_backsolve_bwd(
     # using them later.
     #
 
-    del init_state, t1
+    del perturbed, init_state, t1
     ts, ys = residuals
     del residuals
     grad_final_state, _ = grad_final_state__aux_stats
     # Note that `grad_final_state.save_state` has type `PyTree[SaveState]`; here we are
     # relying on the guard in `BacksolveAdjoint` that it have trivial structure.
     grad_ys = grad_final_state.save_state.ys
+    # We take the simple way out and don't try to handle symbolic zeros.
+    grad_ys = jtu.tree_map(_materialise_none, ys, grad_ys)
     del grad_final_state, grad_final_state__aux_stats
     y, args, terms = y__args__terms
     del y__args__terms
@@ -662,9 +675,6 @@ def __get(__aug):
     return a_y1, a_diff_args1, a_diff_terms1
 
 
-_loop_backsolve.defvjp(_loop_backsolve_fwd, _loop_backsolve_bwd)
-
-
 class BacksolveAdjoint(AbstractAdjoint):
     """Backpropagate through [`diffrax.diffeqsolve`][] by solving the continuous
     adjoint equations backwards-in-time. This is also sometimes known as

diffrax/term.py

@@ -170,7 +170,13 @@ class ODETerm(AbstractTerm):
     vector_field: Callable[[Scalar, PyTree, PyTree], PyTree]
 
     def vf(self, t: Scalar, y: PyTree, args: PyTree) -> PyTree:
-        return self.vector_field(t, y, args)
+        out = self.vector_field(t, y, args)
+        if jtu.tree_structure(out) != jtu.tree_structure(y):
+            raise ValueError(
+                "The vector field inside `ODETerm` must return a pytree with the "
+                "same structure as `y0`."
+            )
+        return jtu.tree_map(lambda o, yi: jnp.broadcast_to(o, jnp.shape(yi)), out, y)
 
     @staticmethod
     def contr(t0: Scalar, t1: Scalar) -> Scalar:

docs/index.md

@@ -20,7 +20,7 @@ _From a technical point of view, the internal structure of the library is pretty
 pip install diffrax
 ```
 
-Requires Python 3.9+, JAX 0.4.4+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.10.4+.
+Requires Python 3.9+, JAX 0.4.13+, and [Equinox](https://github.com/patrick-kidger/equinox) 0.10.8+.
 
 ## Quick example
 

setup.py

@@ -46,7 +46,7 @@
 
 python_requires = "~=3.9"
 
-install_requires = ["jax>=0.4.3", "equinox>=0.10.4"]
+install_requires = ["jax>=0.4.13", "equinox>=0.10.8"]
 
 setuptools.setup(
     name=name,

test/test_adjoint.py

@@ -51,6 +51,17 @@ def _run(y0__args__term, saveat, adjoint):
             ).ys
         )
 
+    def _run_finite_diff(y0__args__term, saveat, adjoint):
+        y0, args, term = y0__args__term
+        y0_a = y0 + jnp.array([1e-5, 0])
+        y0_b = y0 + jnp.array([0, 1e-5])
+        val = _run((y0, args, term), saveat, adjoint)
+        val_a = _run((y0_a, args, term), saveat, adjoint)
+        val_b = _run((y0_b, args, term), saveat, adjoint)
+        out_a = (val_a - val) / 1e-5
+        out_b = (val_b - val) / 1e-5
+        return jnp.stack([out_a, out_b])
+
     diff, nondiff = eqx.partition(y0__args__term, eqx.is_inexact_array)
     _run_grad = eqx.filter_jit(
         jax.grad(
@@ -85,11 +96,15 @@ def _convert_float0(x):
                     continue
                 saveat = diffrax.SaveAt(t0=t0, t1=t1, ts=ts)
 
+                fd_grads = _run_finite_diff(
+                    y0__args__term, saveat, diffrax.RecursiveCheckpointAdjoint()
+                )
                 direct_grads = _run_grad(diff, saveat, diffrax.DirectAdjoint())
                 recursive_grads = _run_grad(
                     diff, saveat, diffrax.RecursiveCheckpointAdjoint()
                 )
                 backsolve_grads = _run_grad(diff, saveat, diffrax.BacksolveAdjoint())
+                assert shaped_allclose(fd_grads, direct_grads[0])
                 assert shaped_allclose(direct_grads, recursive_grads, atol=1e-5)
                 assert shaped_allclose(direct_grads, backsolve_grads, atol=1e-5)
 

test/test_term.py

@@ -36,17 +36,18 @@ def derivative(self, t):
     term = diffrax.ControlTerm(vector_field, control)
     args = getkey()
     dx = term.contr(0, 1)
-    vf = term.vf(0, None, args)
-    vf_prod = term.vf_prod(0, None, args, dx)
+    y = jnp.array([1.0, 2.0, 3.0])
+    vf = term.vf(0, y, args)
+    vf_prod = term.vf_prod(0, y, args, dx)
     assert dx.shape == (2,)
     assert vf.shape == (3, 2)
     assert vf_prod.shape == (3,)
     assert shaped_allclose(vf_prod, term.prod(vf, dx))
 
     term = term.to_ode()
     dt = term.contr(0, 1)
-    vf = term.vf(0, None, args)
-    vf_prod = term.vf_prod(0, None, args, dt)
+    vf = term.vf(0, y, args)
+    vf_prod = term.vf_prod(0, y, args, dt)
     assert vf.shape == (3,)
     assert vf_prod.shape == (3,)
     assert shaped_allclose(vf_prod, term.prod(vf, dt))
@@ -70,17 +71,18 @@ def derivative(self, t):
     term = diffrax.WeaklyDiagonalControlTerm(vector_field, control)
     args = getkey()
     dx = term.contr(0, 1)
-    vf = term.vf(0, None, args)
-    vf_prod = term.vf_prod(0, None, args, dx)
+    y = jnp.array([1.0, 2.0, 3.0])
+    vf = term.vf(0, y, args)
+    vf_prod = term.vf_prod(0, y, args, dx)
     assert dx.shape == (3,)
     assert vf.shape == (3,)
     assert vf_prod.shape == (3,)
     assert shaped_allclose(vf_prod, term.prod(vf, dx))
 
     term = term.to_ode()
     dt = term.contr(0, 1)
-    vf = term.vf(0, None, args)
-    vf_prod = term.vf_prod(0, None, args, dt)
+    vf = term.vf(0, y, args)
+    vf_prod = term.vf_prod(0, y, args, dt)
     assert vf.shape == (3,)
     assert vf_prod.shape == (3,)
     assert shaped_allclose(vf_prod, term.prod(vf, dt))