wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-10 12:14:04 +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/advanced-features.md
+10 -5
View File
@@ -1,13 +1,18 @@
# Advanced features
## Abstract base classes
## Creating your own dtypes
::: jaxtyping.AbstractDtype
selection:
members:
false
::: jaxtyping.AbstractArray
selection:
members:
false
## Introspection
If you're writing your own type hint parser, then you may wish to detect if some Python object is a jaxtyping-provided type.
You can check for dtypes by doing `issubclass(x, AbstractDtype)`. For example, `issubclass(Float32, AbstractDtype)` will pass.
You can check for arrays by doing `issubclass(x, AbstractArray)`. Here, `AbstractArray` is the base class for all shape-and-dtype specified arrays, e.g. it's a base class for `Float32[Array, "foo"]`.
You can check for pytrees by doing `issubclass(x, PyTree)`. For example, `issubclass(PyTree[int], PyTree)` will pass.