Added support for declaring PyTree structures. #135
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
patrick-kidger
force-pushed
the
structure-matching
branch
2 times, most recently
from
50d4a0c
to
6b35a7f
Compare
patrick-kidger
force-pushed
the
fixed-variadic-broadcast
branch
from
c522eb0
to
b661ad0
Compare
patrick-kidger
force-pushed
the
structure-matching
branch
from
6b35a7f
to
380c300
Compare
patrick-kidger
force-pushed
the
fixed-variadic-broadcast
branch
from
b661ad0
to
06feffb
Compare
patrick-kidger
force-pushed
the
structure-matching
branch
2 times, most recently
from
14ffdac
to
d84bd95
Compare
patrick-kidger
force-pushed
the
structure-matching
branch
from
d84bd95
to
0915f9b
Compare
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?)
…his avoids edge-case crash when using pytree-path dependent sizes
…no longer include Float[np.bool, ...].
… from args and kwargs to just arguments.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.