wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Now provides PyTreeDef, and can detect PyTrees via issubclass(x, PyTree)

Browse Source
This commit is contained in:
Patrick Kidger
2023-06-01 10:56:00 -07:00
parent 10e1852b37
commit 6a64ef114e
7 changed files with 68 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/api/pytree.md
+5 -1
View File
@@ -5,4 +5,8 @@
members:
false
Note that `jaxtyping.PyTree` is only available if JAX has been installed.
`jaxtyping.PyTreeDef` is an alias for `jax.tree_util.PyTreeDef`, which is the type of the return from `jax.tree_util.tree_structure(...)`.
:::jaxtyping.PyTreeDef
Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.