Patrick Kidger
Patrick Kidger
Do you have a MWE? (For what it's worth I use generics successfully with Equinox elsewhere.)
Cheers! FWIW I do often also combine Equinox + jaxtyping + beartype. Admittedly that is now a more complicated stack (beartype especially), so I'm definitely willing to believe something goes...
There's been a lot of JAX-specific tricks that have gone into this: for example knowing when the compiler will want to make a copy of a buffer, and avoiding those...
Definitely agreed that the error message could be improved. I'd be happy to take a PR on that! As for performance, you appear to be including the compile time as...
I think let's augment this error message. Tagging @lockwo as I think he may have some idea on this one. Probably we should write out something quite verbose -- in...
I think probably easiest to do these as options, as long as the branching isn't too complicated! (The main bottleneck is my ability to review/maintain whatever is added.) FWIW whilst...
Another +1 for this. Setting up fonts is always one of the most fiddly things in any system (Typst, Latex, Matplotlib, ...). This is something that I'd like to be...
I have a hunch that this is actually a JAX bug: namely, that it is spuriously introducing a batch tracer (they do this in a few places where it made...
> Your example shows an error that recommends using eqx.filter_grad over jax.grad, have you tried this? I think you might be misreading the error message, this is one recommending using...
I'm afraid there isn't. The intended interpretation of this is to check that the shapes are the same.