From 29ecb3b5d0b1851d49834bf4e6d901738324e645 Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 18 Aug 2026 22:37:10 -0400 Subject: [PATCH] allocate the SaveAt buffer from the saved value, not its shape `jnp.full` conjures an array with no operand, so an array-ish abstraction wrapping `y0` -- a `quax.Value` carrying units, an uncertainty, a sparsity pattern -- has nothing to dispatch on and is erased as soon as a saved value round-trips through the buffer via the `lax.cond` in `_save`. Route the same `inf` fill through a select whose predicate is a compile-time constant, so `y` is an operand. XLA folds the select away and DCEs the dead branch, including the `subsaveat.fn` call feeding it: the cost analysis is unchanged (4 fewer flops, identical bytes accessed and temp memory) and saved values, including the `inf` fill of unwritten slots, are bit-identical. The `stop_gradient` is load-bearing -- without it the buffer acquires a zero-valued but structurally present tangent path back to `y0`, tripping `BacksolveAdjoint`'s `nondifferentiable` guard. Co-Authored-By: Claude Opus 5 --- diffrax/_integrate.py | 19 +++++++-- pyproject.toml | 1 + test/test_quax.py | 89 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 4 deletions(-) create mode 100644 test/test_quax.py diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py index 5441e0a9..99ae21e9 100644 --- a/diffrax/_integrate.py +++ b/diffrax/_integrate.py @@ -1294,10 +1294,21 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState: saveat_ts_index = 0 save_index = 0 ts = jnp.full(out_size, direction * jnp.inf, dtype=time_dtype) - struct = eqx.filter_eval_shape(subsaveat.fn, t0, y0, args) - ys = jtu.tree_map( - lambda y: jnp.full((out_size,) + y.shape, jnp.inf, dtype=y.dtype), struct - ) + + # Allocate the buffer *via* an operation on the saved value, not from its + # shape: `jnp.full` conjures an array with no operand, so an array-ish + # abstraction wrapping `y0` (a `quax.Value`) is erased as soon as a saved + # value round-trips through it (`_save` slices it opposite `y` in a `cond`). + # The select's predicate is a compile-time constant, so XLA folds it away + # and DCEs the `subsaveat.fn` call feeding the dead branch. `stop_gradient` + # is required: without it the buffer gains a tangent path back to `y0` and + # `BacksolveAdjoint` trips its `nondifferentiable` guard. + def _alloc(y): + shape = (out_size,) + jnp.shape(y) + fill = jnp.full(shape, jnp.inf, dtype=jnp.result_type(y)) + return jnp.where(True, fill, lax.stop_gradient(jnp.broadcast_to(y, shape))) + + ys = jtu.tree_map(_alloc, subsaveat.fn(t0, y0, args)) return SaveState( ts=ts, ys=ys, save_index=save_index, saveat_ts_index=saveat_ts_index ) diff --git a/pyproject.toml b/pyproject.toml index 51cfe6d4..6ae595f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ tests = [ "jaxlib", "optax", "pytest", + "quax", "scipy", "tqdm" ] diff --git a/test/test_quax.py b/test/test_quax.py new file mode 100644 index 00000000..d6fd06d2 --- /dev/null +++ b/test/test_quax.py @@ -0,0 +1,89 @@ +"""`diffeqsolve` under `quax.quaxify`. + +Quax dispatches on *operations*. An array conjured from a shape rather than +from an operand -- `jnp.full(shape, ...)` -- gives it nothing to dispatch on, +so a wrapped `y0` is silently erased the moment a saved value round-trips +through the `SaveAt` buffer. These tests pin the buffer as an operation on +the value being saved. +""" + +from collections.abc import Sequence + +import diffrax +import equinox as eqx +import jax +import jax.numpy as jnp +import pytest +from jax.core import ShapedArray +from jax.extend.core import Primitive + + +# `quaxify(diffeqsolve)` additionally needs quax's `custom_vjp` support, which +# diffrax's buffered loops go through; that landed in quax 0.5.0. +quax = pytest.importorskip("quax", minversion="0.5.0") + + +class Boxed(quax.ArrayValue): + """Propagates itself through any primitive it is an operand of. + + Boxes *inexact* outputs only. Boxing integers too would put the box on the + solver's own loop counters and save indices, at which point essentially + every primitive has a boxed operand and the buffer gets re-boxed on its + first write -- which hides exactly the bug under test (and is why + `sol.ts` came back wrapped in diffrax#438). + """ + + array: jax.Array = eqx.field(converter=jnp.asarray) + + def aval(self) -> ShapedArray: + return ShapedArray(jnp.shape(self.array), jnp.result_type(self.array)) + + def materialise(self) -> jax.Array: + return self.array + + @staticmethod + def default(primitive: Primitive, values: Sequence, params: dict): + raw = [v.array if isinstance(v, Boxed) else v for v in values] + out = primitive.bind(*raw, **params) + + def box(x): + inexact = eqx.is_array(x) and jnp.issubdtype(x.dtype, jnp.inexact) + return Boxed(x) if inexact else x + + return [box(x) for x in out] if primitive.multiple_results else box(out) + + +def _solve(y0, saveat): + return diffrax.diffeqsolve( + diffrax.ODETerm(lambda t, y, args: -0.5 * y), + diffrax.Euler(), + t0=0.0, + t1=1.0, + dt0=0.1, + y0=y0, + saveat=saveat, + # `throw=True` erases the type again via equinox's `error_if`: + # patrick-kidger/equinox#1257. Not a diffrax issue. + throw=False, + ) + + +@pytest.mark.parametrize( + "saveat", + [ + diffrax.SaveAt(t1=True), + diffrax.SaveAt(t0=True, t1=True), + diffrax.SaveAt(ts=[0.0, 0.5, 1.0]), + diffrax.SaveAt(steps=True), + ], + ids=["t1", "t0t1", "ts", "steps"], +) +def test_saveat_buffer_preserves_quax_type(saveat): + y0 = jnp.array([1.0]) + expected = _solve(y0, saveat).ys + assert expected is not None + + got = quax.quaxify(_solve)(Boxed(y0), saveat).ys + + assert isinstance(got, Boxed), f"`sol.ys` came back as {type(got).__name__}" + assert jnp.array_equal(got.array, expected, equal_nan=True)