Owen L
Owen L
This stemmed from the issues of allowing non array inputs and outputs, and just shifted to only allowing non differentiable inputs (see more: https://github.com/patrick-kidger/equinox/pull/734). If you have no static outputs...
In the end the alternative approach was what we went with (since it had precedent in equinox, and also came up in some discussions on jax issues with the core...
> If you want Davie approximations to Levy area then I think we had that in torchsde. It does exist! This torchsde package sounds pretty neat, I might have to...
> Okay, I really like the example here! > > I'm afraid this might still take a bit more iteration to get to something clean enough to merge, though --...
Looks like the issue still persists, just checked on firefox 
Ok, I took the feedback from Andraz's Langevin PR regarding terms and incorporated it into this PR. I think it made things simpler and more in line with the diffrax...
There's a fair amount that could be said here, so first I'll just remark that with certain diffrax parameters you can come much closer to the lax.scan time. Specifically, by...
I am not aware of any examples using TFP + equinox, but it should be possible given that equinox generally operates at the jax level, rather than as a wrapper....
To do what you're describing there's a couple ways (since I assume you don't much care about that specific problem, but want to apply it to your, potentially much more...
The gradients is a pytree of the same structure as the parameters, so wherever in the parameter pytree the bias is, same in the gradients