Patrick Kidger

Hmm, is this not just the overhead from doing the actual type checking itself? For what it's worth I don't think we currently respect `@typing.no_type_check` for dataclasses (only functions), but...

Oh interesting! Thank you for profiling this -- I agree, caching sounds reasonable. I'd be happy to take a PR on this!

So this is expected, as `filter_pmap` actually uses an additional `vmap` call under the hood: https://github.com/patrick-kidger/equinox/blob/36f81165d53328a9e83fc1dc67afc0dc25fa48ec/equinox/_vmap_pmap.py#L372 in order to resolve the output axes correctly. I can see the interaction with...

It's pretty edge-case, unless this comes up a few times then I think just the existence of this issue is probably enough :)

I worry that privileging *just* single-cell functions is going to be a huge UX footgun. There are many other things I might reasonably want the same behaviour from. (Classes are...

Okay, refining this idea a little bit: it should be possible to define a module-level `__getattr__` that runs the appropriate cells up to the point of providing that name. That...

Agreed, this sounds like a use case for [writing a custom solver](https://docs.kidger.site/diffrax/usage/extending/). I'd recommend taking a look at the source code for a simple solver like [Euler](https://github.com/patrick-kidger/diffrax/blob/main/diffrax/_solver/euler.py) for an example...

Ah, sorry -- I didn't understand that detail. So I think you could still make this happen with a custom solver. A solver gets to return a `result` specifying whether...

I'd be happy to add these to the public API. Just never had a request for that before! On the topic of functional APIs, one of the nice thinks about...

Should this be something like `jnp.finfo(jnp.float32).max`? If nothing else to document the choice here.