Patrick Kidger
Patrick Kidger
The crash with `Shaped[A | B, ...]` should be fixed in #77 (version 0.2.15).
`A | B` should already be supported, as per #77. You appear to be linking to a commit from about 2 years ago, prior to this fix :)
I suspect you also need to set `EQX_ON_ERROR_BREAKPOINT_FRAMES`; see also the docs at https://docs.kidger.site/equinox/api/errors/#equinox.error_if. This is all to work around [a known JAX bug](https://github.com/google/jax/issues/16732), which is likely what you're bumping...
So the stack goes function -> jaxpr -> runtime. I think if the trace-time callstack doesn't help you, then I do have some tricks up my sleeve to help with...
Ah, that's good to know! For what it's worth `from __future__ import annotations` will often break runtime type checkers in unpredictable ways (resolving the stringified annotations correctly is intractable in...
Indeed, callable nodes are possible, so for this reason it has to be a dynamic attribute. :)
Sounds reasonable to me! Happy to take a PR on this.
Interesting point! Supportig this is definitely a non-goal for jaxtyping, although I'd be happy to accept PRs if something needs tweaking to enable torchscript compatibility. For what it's worth, my...
Thanks for the suggestion! The libraries I recommend here are actually an intentionally short list: for the most part they are ones that I've used myself. However, I can totally...
I think this is happening because you're grabbing `__call__`, which as a magic method isn't subject to the same bound-methods-are-PyTrees treatment as regular methods. This is the reason `t` is...