mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-09-09 11:24:55 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1456302503 | ||
|
|
1c0186a1c9 |
@@ -1,6 +1,6 @@
|
|||||||
<h1 align="center">jaxtyping</h1>
|
<h1 align="center">jaxtyping</h1>
|
||||||
|
|
||||||
Type annotations **and runtime type-checking** for:
|
[Use type annotations **and runtime type-checking**](https://jax.readthedocs.io/en/latest/jep/12049-type-annotations.html) for:
|
||||||
|
|
||||||
1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
|
1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
|
||||||
2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
|
2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
|
||||||
@@ -8,7 +8,11 @@ Type annotations **and runtime type-checking** for:
|
|||||||
|
|
||||||
**For example:**
|
**For example:**
|
||||||
```python
|
```python
|
||||||
from jaxtyping import Array, Float, PyTree
|
from jaxtyping import Array, Float, PyTree, , UInt, Int, Bool
|
||||||
|
import torch
|
||||||
|
impport numpy as np
|
||||||
|
import tensorflow as tf
|
||||||
|
|
||||||
|
|
||||||
# Accepts floating-point 2D arrays with matching axes
|
# Accepts floating-point 2D arrays with matching axes
|
||||||
def matrix_multiply(x: Float[Array, "dim1 dim2"],
|
def matrix_multiply(x: Float[Array, "dim1 dim2"],
|
||||||
@@ -21,6 +25,15 @@ def accepts_pytree_of_ints(x: PyTree[int]):
|
|||||||
|
|
||||||
def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
|
def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
|
||||||
...
|
...
|
||||||
|
|
||||||
|
def accepts_torch.Long(x: Int[torch.Tensor, "batch channel height width"]):
|
||||||
|
....
|
||||||
|
|
||||||
|
def accepts_numpy_float(x :Float[np.ndarray, "batch sequence features"]):
|
||||||
|
...
|
||||||
|
|
||||||
|
def accepts_tensorflow_uint(x: hint = UInt[tf.Tensor, "b c h w"]):
|
||||||
|
...
|
||||||
```
|
```
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|||||||
Reference in New Issue
Block a user