mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-25 11:19:16 +08:00
Patrick Kidger
12d540794f
Phew, this ended up being a pretty complicated change! The basic summary is that we now support the syntax ``` @jaxtyped(typechecker=typechecker) def f(...): ... ``` and when using this, we now get pretty error messages about what went wrong. ( The old syntax, i.e. ``` @jaxtyped @typechecker def f(...): ... ``` is still supported, but doesn't give much information. ) The internals of this do quite a lot of magic! In particular we dynamically create quite a lot of functions and test the provided arguments against their signatures. The overhead should still be minimal under `jax.jit`, though. (TODO: what's the overhead like in non-jit situations, e.g. PyTorch? I've tried to minimise the overhead throughout just to be sure, but perhaps PyTorch users should stick to the old syntax?)