Patrick Kidger

I'd be happy to offer this if you can determine an API for it. Because this happens inside a `custom_vjp` then as far as I can see JAX doesn't offer...

I probably wouldn't mutate using `pop`, but other than that it LGTM. Something else that might make logging easier is using `equinox.experimental.{get,set}_state` instead of `host_callback` directly. See [here](https://docs.kidger.site/equinox/api/experimental/stateful/). This wraps...

This is very elegantly written. On balance I'm still inclined to host these in an external project, but as before I'm very happy to advertise its existence. What I'd quite...

I'm closing this as we now have `paganpasta/eqxvision`.

Doing so with MLIR: right! After implementing the above I started wondering something similar. As an end goal, could we arrange to generate a single `XlaComputation` that all usages point...

> If wrap with rng fake while loop can successfully prevent inlining(really?), implement `noinline` should be easier. But I guess this trick can only prevent computation from moving outside. XLA...

Closing in favour of #218, which adds [`equinox.internal.noinline`](https://github.com/patrick-kidger/equinox/blob/1f5373f5905504bc5e7069ed6d458dbad5616495/equinox/internal/noinline.py#L356). This offers a heavily improved version of this. - Much faster. - Uses a custom primitive directly, instead of monkey-patching `host_callback.call`. -...

A standard approach for something like this would be a [shooting method](https://en.wikipedia.org/wiki/Shooting_method).

Hmm. Unless there's a big outcry from lots of people wanting this feature, then I'm inclined to keep things as-is just for simplicity. However I expect it should be possible...

Hi there; I'm glad you like the library. So the answer is that handling the general case for `...` is complicated. For example someone might want to match against `TensorType[...,...