repository·main·Indexed 22 days ago
https://github.com/patrick-kidger/jaxtypingType annotations and runtime checking for the shape and dtype of arrays and tensors across multiple frameworks, including JAX, PyTorch, NumPy, MLX, and TensorFlow. It allows specifying data types and dimensions using string literals and supports integration with runtime checkers like typeguard and beartype, as well as JAX-specific PyTree annotations.
jax.jit. When using jax.jit, dtype and shape checking occurs at trace time (when JAX traces the function prior to compilation). The resulting compiled code does not contain any dtype/shape-checking overhead, ensuring performance remains identical to untyped JAX code.You can annotate arrays using the syntax dtype[array_type, shape]. This allows you to specify both the expected data type and the dimensions of the array in a single type hint.
Example:
jaxtyping.Float[torch.Tensor, "batch channels"] specifies a PyTorch tensor where the data is floating-point and the shape has two axes named batch and channels.
jaxtyping.Float[torch.Tensor, "batch channels"]You can use the ? prefix in a shape string (e.g., ?foo) to indicate that an axis size is path-dependent. This means the size of that axis can vary depending on which leaf of the PyTree is being inspected, but it must remain consistent across different PyTrees sharing the same structure template.
When you use ?foo, jaxtyping internally treats each leaf as having a unique version of that dimension (e.g., 0foo, 1foo, etc.). This allows you to enforce that corresponding leaves in two different PyTrees have matching dimensions, even if the dimensions differ between leaves within the same tree.
def f(
x: PyTree[Shaped[jax.Array, "?foo"], "T"],
y: PyTree[Shaped[jax.Array, "?foo"], "T"],
):
pass
# Example of valid usage:
x0 = jnp.arange(3)
x1 = jnp.arange(5)
y0 = jnp.arange(3) + 1
y1 = jnp.arange(5) + 1
# This works because x0 matches y0 and x1 matches y1
f((x0, x1), (y0, y1))
# This fails because x1 (size 5) does not match y0 (size 3)
# f((x1, x1), (y0, y1))jaxtyping provides partial support for static type checkers. Annotations like dtype[array, shape] are treated as just array by static checkers, as full dtype/shape checking is currently beyond the scope of static type checking capabilities.
dtype[array, shape] is sometimes treated as Any instead of array.The following frameworks are supported as the array_type argument:
jax.Array / jax.numpy.ndarraynp.ndarray (NumPy)torch.Tensor (PyTorch)tf.Tensor (TensorFlow)mx.array (MLX)Duck-typing: You can use any object that has .shape (returning tuple[int, ...]) and .dtype (returning a str). To use custom dtypes with duck-typed arrays, inherit from jaxtyping.AbstractDtype and define the dtypes list.
class MyDuckArray:
@property
def shape(self) -> tuple[int, ...]:
return (3, 4, 5)
@property
def dtype(self) -> str:
return "my_dtype"
class MyDtype(jaxtyping.AbstractDtype):
dtypes = ["my_dtype"]
x = MyDuckArray()
assert isinstance(x, MyDtype[MyDuckArray, "3 4 5"])Shapes are defined as space-separated strings. Supported symbols include:
int: A fixed-size axis (e.g., "28 28").str: A variable-size axis name (e.g., "channels"). These names are used to match axes across different function arguments."dim-1"). These must not contain spaces....: Anonymous zero or more axes (equivalent to *_)."" (Empty string): Represents a scalar shape.Note on Symbolic Expressions: To use local variables or class attributes in symbolic expressions, wrap them in curly braces {}. They are evaluated as f-strings first, then as axis sizes.
def full(size: int, fill: float) -> Float[jax.Array, "{size}"]:
return jax.numpy.full((size,), fill)
class SomeClass:
some_value = 5
def full(self, fill: float) -> Float[jax.Array, "{self.some_value}+3"]:
return jax.numpy.full((self.some_value + 3,), fill)jax.jit. All shape checks are performed during the tracing process and do not impact the runtime performance of the compiled function.If you are building tools like type hint parsers or decorators, you can detect if a Python object is a jaxtyping type using issubclass.
issubclass(x, AbstractDtype).issubclass(x, AbstractArray).issubclass(x, PyTree).Install jaxtyping using pip. Note that it requires Python 3.10 or higher.
To enable runtime type-checking of array shapes and dtypes, it is recommended to also install a runtime type-checking library such as typeguard or beartype.
pip install jaxtypingIn IPython environments like Jupyter or Colab, you can use a magic command to automatically apply the jaxtyping hook to everything defined in the notebook after the magic is run.
import jaxtyping
%load_ext jaxtyping
%jaxtyping.typechecker beartype.beartype # or any other runtime type checker