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

Added support for jax.typing.ArrayLike; now works with PyTorch's bool

Browse Source
This commit is contained in:
Patrick Kidger
2023-02-25 17:01:27 -08:00
parent 81c56052e5
commit e2f004afd4
6 changed files with 222 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4 -6
View File
@@ -29,6 +29,10 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
pip install jaxtyping
```
Requires Python 3.8+.
JAX is an optional dependency, required for `jaxtyping.{Array, ArrayLike, PyTree}`. If JAX is not installed then these types will not be available, but you may still use jaxtyping alongside PyTorch/NumPy/etc.
Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
## Documentation
@@ -49,12 +53,6 @@ Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
### Acknowledgements
Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping).
The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types).
### Disclaimer
This is not an official Google product.