Fix for https://github.com/patrick-kidger/optimistix/issues/48

[patrick-kidger]

This is a fix for patrick-kidger/optimistix#48

It happens occasionally that JAX will insert a spurious vmap. It seems like we're hitting one of these cases.

In particular I found that:

• We are hitting this nonbatchable call

  equinox/equinox/internal/_loop/common.py

  Line 432 in 60612c1

  return nonbatchable(out)

  , due to step getting batched.
• Inserting a val = (nonbatchable(val[0]),) + val[1:] here:

  equinox/equinox/internal/_loop/bounded.py

  Line 72 in 60612c1

  def scan_fn(val, _):

  will also trigger an error before the previously-mentioned one. But shifting this additional check to the other side of the jax.checkpoint (a few lines further down) will not raise an error before the previously-mentioned one.

Spelunking through the (very deep) JAX stack traces, it seems to be due to the lax.scan saving its evolving state for use in the backward pass -- in particular including this step -- and for some reason the insertion of the jax.checkpoint causes this state to become batched. Then when we run the backwards scan during backprop, we re-run this function with the batched state and so things explode. I didn't try digging into it further than that, as I'm not hugely surprised -- such spurious batching is a thing JAX does every now and again.

In terms of a fix: for the actual mathematical correctness, we do continue to disallow a batch to have nonconstant steps (I would need to think a lot harder about whether nonconstant step can ever be valid). However if a batch tracer does arise, then we simply punt the nonconstancy check to runtime. This definitely isn't ideal, but this is (a) a pretty unusual edge-case (DirectAdjoint + vmap + optimistix), and (b) it's better than crashing. (Hopefully all of this becomes unnecessary if jvp-of-custom_vjp is implemented and we can remove DirectAdjoint entirely.)