Patrick Kidger

Levenberg-Marquardt is already one of the possible solvers we have available. What you are observing is that, in addition to that, we also support converting the residuals into a minimization...

torchsde is in in maintenance mode only, so I now recommend Diffrax. See the 'see also' section of the last torchsde release: https://github.com/google-research/torchsde/releases/tag/v0.2.6 On a technical level, Diffrax is typically...

Swap your final line for: ```python print(jax.tree.map(lambda x, f : f(x), tst_states, tst_fxns_eqx)) ``` This isn't a bug. Equinox modules are pytrees. `jax.tree.map` requires that the first tree (in your...

Happy to do so! For now just waiting on the JAX 0.7.2 release as both JAX 0.7.0 and JAX 0.7.1 have some incompatibilities/regressions with the Equinox ecosystem which should hopefully...

> Happy to address these during the week, just wanted to check: I did a quick trawl through the code and I don't think diagonal or tridagonal is ever used...

Sorry for the long delay getting back to you, some personal life things took over for a while. So, now to actually answer your question: good point. I imagine we...

This is probably a variant of https://github.com/jax-ml/jax/issues/16732. Equinox already has a workaround for the specific reported version above (when `EQX_ON_ERROR=breakpoint` is set then we monkey-patch `jax.jit` to conditionally disable it),...

This is probably possible with a custom solver. For some initial familiarity on how they work, we have some docs offering the starting point for doing so [here](https://docs.kidger.site/optimistix/abstract/). After that...

I think probably something like this is what you're looking for: ```python @eqx.filter_jit(donate="all-except-first") def _do_stuff(aux, x, example): return example.loads_of_logic(x, aux) class Example(eqx.Module): def do_stuff(self, x, aux): return _do_stuff(aux, x, self)...

So this isn't something we have a nice way to annotate in jaxtyping right now. I think supporting this is the same as https://github.com/patrick-kidger/jaxtyping/issues/242. I'm not sure how fiddly that...