Feature request: Support python control flow in `custom_transforms` functions

For fitting parameter values for ODEs a la the adjoint sensitivity method, we might want to override the gradient computation for the forward ODE solve. More concretely, we might have an integrator function odeint that takes in a gradient field f, initial state y0, and a sequence of times ts to be evaluated at.

One specific use case where supporting control-flow in custom_transforms will be useful is for the backward integration (which might involve adaptive solvers, hence non-trivial control-flow). Ideally, we would like to write code as follows

@custom_transforms
def odeint(y0, ts):
  pass  # Some procedure integrating the vector field `f`.

def vjp_y0(g, ans, y0, ts):
  pass # A while loop and some if statements used to determine integration step size.

defvjp(odeint, vjp_y0, None)

[mattjj]

We’re working on this. The reason it isn’t as simple as it was for Autograd is that JAX uses a new autodiff design in which we only have forward mode and derive reverse mode automatically (composing forward mode with other transformations). That confers several advantages, but a disadvantage is that since the system itself doesn’t work in terms of VJPs, supporting custom VJPs is tricky. (You can write custom JVPs, ie forward-mode rules, with arbitrary Python control flow now.)

[mattjj]

#2026 finally landed and added support for Python control flow in custom derivative rules! (The API also changed, so take a look at the tutorial notebook.)