Allocate the SaveAt buffer from the saved value, not its shape

[nstarman]

Allocates the SaveAt output buffer via an operation on the value being saved, rather than conjuring it from that value's shape and dtype.

The problem

_allocate_output builds the buffer with jnp.full:

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)

jnp.full takes no operand. So when y0 is an array-ish abstraction — a quax.Value carrying units, an uncertainty, a sparsity pattern — there is nothing here for Quax to dispatch on, and the buffer comes out as a plain array. _save then puts a slice of that plain buffer opposite the carried value in a lax.cond:

y_to_save = lax.cond(
    pred,
    lambda: fn(t, y, args),                                  # the wrapped value
    lambda: jtu.tree_map(lambda ys_: ys_[save_index], ys),   # a plain buffer slice
)

The two branches disagree on structure, cond_p needs one, and the wrapper is dropped. sol.ys comes back as a bare array; a type that declines to materialise (units, say — a dimensionless metre has no meaning) raises instead.

Why this one can't be fixed from outside

In #438 the conclusion was that Quax shouldn't require changing Diffrax, and for essentially all of diffeqsolve that holds — Quax intercepts primitives, and Diffrax is written in primitives. This site is the exception, and structurally so: a rule can only fire where there is an operand to dispatch on, and an array conjured from a shape has none. No set of Quax rules can reach it.

(The workaround in #438 — a default rule permissive enough to re-box everything — only appears to work because it also boxes the solver's integer loop counters, at which point almost every primitive has a boxed operand and the buffer gets re-boxed on first write. That is also why sol.ts came back wrapped there.)

The change

Keep the inf fill exactly, but route it through a select whose predicate is a compile-time constant, so y is an operand:

y_init = subsaveat.fn(t0, y0, args)
ys = jtu.tree_map(
    lambda y: jnp.where(
        True,
        jnp.full((out_size,) + jnp.shape(y), jnp.inf, dtype=jnp.result_type(y)),
        lax.stop_gradient(jnp.broadcast_to(y, (out_size,) + jnp.shape(y))),
    ),
    y_init,
)

XLA folds the select away and DCEs the dead branch — including the subsaveat.fn call feeding it.

The stop_gradient is load-bearing. Without it the buffer acquires a tangent path back to y0 — zero-valued, but structurally present — and BacksolveAdjoint trips its own nondifferentiable guard: RuntimeError: Unexpected tangent. That showed up as 4 failures in test_adjoint.py. With stop_gradient all adjoints produce gradients identical to main.

Cost

XLA cost analysis, 64-dim state, Tsit5 + PIDController(rtol=1e-6, atol=1e-8), max_steps=4096:

SaveAt	flops before → after	bytes accessed	temp memory
t1=True	3,267 → 3,263	14,113 → 14,113	3,392 → 3,392
steps=True	3,440 → 3,436	1,080,575 → 1,080,575	1,065,888 → 1,065,888
steps=True, fn=sin(y)**3+cos(y)	3,632 → 3,628	1,080,575 → 1,080,575	1,065,888 → 1,065,888

Four fewer flops, identical memory traffic. The non-trivial fn row is there to show the extra subsaveat.fn call costs nothing once the dead branch is eliminated.

Behaviour

Unchanged. Saved values are bit-identical and unwritten slots still read back as inf — including test_saveat_solution.py's explicit jnp.allclose(ys1, jnp.array([0.25, 0.09375, jnp.inf])).

One caveat worth naming: subsaveat.fn is now invoked once during allocation. It was already invoked inside the solve, so this introduces no new kind of behaviour, but a fn with a side effect (a jax.debug.print, say) would fire one additional time, since DCE removes the pure computation but not the effect.

Two simpler forms that I tried and discarded, in case they come to mind on review:

• jnp.broadcast_to(y_init, ...) with no select at all. Changes the fill from inf to y0, which contradicts the assertions above, and separately broke test_adjoint.py::test_against.
• The select without the stop_gradient, which breaks BacksolveAdjoint as described.

Test

test/test_quax.py asserts sol.ys keeps its wrapper across t1 / t0,t1 / ts / steps, that sol.ts does not pick one up, and that the inf fill survives. Verified against this branch: 4 failures before the change, all green after.

It is gated on quax>=0.5.0 and so currently skips — quaxify(diffeqsolve) also needs Quax's custom_vjp support (Diffrax's buffered loops route through jax.custom_vjp), which lands in that release. Happy to drop the test and the test-group dependency if you would rather not carry a skipped test until then.

Not fixed here

With the default throw=True, the Solution is passed through equinox.internal.error_if, which allocates its error branch from jax.eval_shape + pure_callback — the same shape-not-operand pattern, one library over. That erases the type again, so the test uses throw=False. Filed as patrick-kidger/equinox#1257, with a candidate patch.

[patrick-kidger]

Sorry for the delay, gosh.

Just so that I understand this and patrick-kidger/equinox#1257 correctly: is your goal to introduce a lax.select_n_p.bind(predicate, raw_array, quax_value), which then rules like

https://github.com/nstarman/quax/blob/9ae62578fc877c7098c068d1d316982eea2346f4/src/quax/examples/unitful/_core.py#L125-L126

can catch on? And then you write your Quax rule to be okay with receiving an array/quax combo, and to create its output as a Quax Value? Let me know if I've understood that correctly or if I'm misunderstanding.

[nstarman]

Yes, that's the gist.
As a note, jnp.where(True, fill, y) lowers to select_n(pred, y, fill), so the Value is case 0 and the buffer is case 1. And we also need to dispatch for broadcast_in_dim and stop_gradient.

On whether rules can catch it: the linked select_n_unitful(which, *cases: Unitful) doesn't match (Unitful, Array), so it falls to Value.default -> materialise -> Unitful raises. With these two, the alloc comes back as Unitful(inf-buffer, {m: 1}):

@quax.register(jax.lax.stop_gradient_p)
def _(x: Unitful, **kw) -> Unitful:
    return Unitful(jax.lax.stop_gradient_p.bind(x.array, **kw), x.units)

@quax.register(jax.lax.select_n_p)
def _(which: ArrayLike, case0: Unitful, case1: ArrayLike, **kw) -> Unitful:
    return Unitful(jax.lax.select_n_p.bind(which, case0.array, case1, **kw), case0.units)

So yeah it's array/Value combo in, Value out.

The part I'm least happy about is that these dispatches are ill-typed on its face. What we actually want is a construction but quax dispatches on operations only, so this isn't possible in diffrax without some creative use of map operations. Currently, a select with a constant predicate is the smallest way to set up a construction as an operation. An alternative is a full_like hook quax side, which I think diffrax would have to use.