wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-08-28 12:50:23 +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
API.md
+2
View File
@@ -62,6 +62,8 @@ Float32[Array, "some_shape"]
The array should typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`.
`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`.
But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.:
```python
Float[np.ndarray, "..."]