Archis Joglekar

Perhaps this should go to a different issue but I guess I'll start here. Since we're talking about launching CUDA/cuDNN kernels, is there any desire in implementing a wrapper to...

Makes sense. Glad I wasn't totally off-base. Thanks for the prompt response, @mattjj

> > The most straightforward way would be to include the forcing term in vector_field - if one had a functional representation of it. Let's say that the forcing signal...

> Do I understand that: > > 1. As in [Logging metrics during an ODE solve #60](https://github.com/patrick-kidger/diffrax/issues/60), you're interested in logging metrics during the solve, and want to save something...

That's exactly what I do right now in the scan (in a Haiku module, I know, I know, transform magic, i'm certainly thinking about equinox) ```python def one_step(y, current_params): y...

i am in, is this where we'll be working? https://github.com/patrick-kidger/diffrax/blob/cec091c5e4cc4311f64ae3aa09a371db5fe766ee/diffrax/integrate.py#L246

Looking at this again, do I want to trace back the `ys=ys` line to be more of a `out=out` line? ```python new_state = _State( y=y, tprev=tprev, tnext=tnext, made_jump=made_jump, solver_state=solver_state, controller_state=controller_state,...

Thanks for helping me get this through. `SubSaveAt` is a nice feature add too! I think I'm on to a bug, but I don't have a minimum repro for this...

I think I have this traced down to storing `abs` and `angle` of complex values. This happens even if I take the gradient of a different quantity that never undergoes...

I don't think this is fully related but since you mentioned `hessian` here, I'm running into `TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.` for a simple Van...