### Clone and Install Jaxtyping

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/CONTRIBUTING.md

Clone the repository, navigate to the directory, and install dependencies including pre-commit hooks using uv.

```bash
git clone https://github.com/your-username-here/jaxtyping.git
cd jaxtyping
uv run prek install  # Creates a local venv + installs dependencies + installs pre-commit hooks.
```

--------------------------------

### Install Jaxtyping

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/README.md

Install the jaxtyping library using pip. Ensure you have Python 3.10 or newer.

```bash
pip install jaxtyping
```

--------------------------------

### Example of Symbolic Expression in Return Type

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Demonstrates using symbolic expressions in return type annotations, referencing axes declared in arguments. This allows for dynamic shape calculations.

```python
def remove_last(x: Float[torch.Tensor, "dim"]) -> Float[torch.Tensor, "dim-1"]:
    pass
```

--------------------------------

### JAX-Specific Type Aliases and Usage

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Demonstrates the use of JAX-specific type aliases like Array, Scalar, and PRNGKeyArray with jaxtyped and beartype for runtime checking. Shows examples of function definitions and usage with JAX arrays.

```python
import jax
import jax.numpy as jnp
from jaxtyping import Array, ArrayLike, Scalar, PRNGKeyArray, Key, Float, jaxtyped
from beartype import beartype

@jaxtyped(typechecker=beartype)
def safe_norm(x: Float[Array, "n"]) -> Scalar:
    return jnp.linalg.norm(x)

@jaxtyped(typechecker=beartype)
def split_key(key: PRNGKeyArray) -> Float[Array, "2 ..."]:
    return jax.random.split(key, 2)

@jaxtyped(typechecker=beartype)
def accepts_arraylike(x: Float[ArrayLike, "batch dim"]) -> Array:
    return jnp.asarray(x)

key = jax.random.key(42)
print(safe_norm(jnp.array([3.0, 4.0])))   # 5.0
print(split_key(key).shape)               # (2,) — new-style keys are scalar

# Composed annotations
from jaxtyping import Shaped
BatchPRNG = Shaped[PRNGKeyArray, "n"]
print(isinstance(jax.random.split(key, 4), BatchPRNG))  # True
```

--------------------------------

### Printing Axis Bindings

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/advanced-features.md

Utilize the print_bindings function to display axis bindings within your Jaxtyping setup.

```APIDOC
## Printing axis bindings

::: jaxtyping.print_bindings
```

--------------------------------

### Duck-type Array Annotation Example

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Provides an example of annotating a duck-type array that has .shape and .dtype attributes. This demonstrates how Jaxtyping can work with custom array-like objects.

```python
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"])
# checks that `type(x) == MyDuckArray`
# and that `x.shape == (3, 4, 5)`
```

--------------------------------

### Dataclass Annotation Example

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/faq.md

Demonstrates how stringified dataclass annotations are silently skipped by Jaxtyping, and partially stringified annotations will likely raise an error. Avoid using stringified annotations with `from __future__ import annotations`.

```python
from dataclasses import dataclass

@dataclass()
class Foo:
   x: "int"
```

```python
from dataclasses import dataclass

@dataclass()
class Foo:
   x: tuple["int"]
```

--------------------------------

### Equivalent BatchImage type definition

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Shows the expanded form of the BatchImage type alias, illustrating how nested type specifications are resolved. This is equivalent to the previous example but explicitly shows the flattened shape.

```python
BatchImage = Float[jax.Array, "batch channels height width"]
```

--------------------------------

### Define Typed Function with Jaxtyping

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/README.md

Use Jaxtyping annotations to define function signatures that enforce specific tensor shapes and dtypes at runtime. This example shows a matrix multiplication function.

```python
from jaxtyping import Float
from torch import Tensor

# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Tensor, "dim1 dim2"],
                    y: Float[Tensor, "dim2 dim3"])
                  -> Float[Tensor, "dim1 dim3"]:
    ... 
```

--------------------------------

### Serve Documentation Locally

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/CONTRIBUTING.md

Build and serve the documentation locally using mkdocs to preview changes. Access it at localhost:8000.

```bash
uv run mkdocs serve
```

--------------------------------

### Shape Annotation Syntax with PyTorch

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Illustrates shape annotation syntax including named axes, fixed-size integers, variadic axes (*), broadcasting axes (#), and scalar arrays. Uses PyTorch tensors and beartype for checking.

```python
import torch
from jaxtyping import Float, jaxtyped
from beartype import beartype

@jaxtyped(typechecker=beartype)
def batch_outer(x: Float[torch.Tensor, "b c1"],
                y: Float[torch.Tensor, "b c2"]
                ) -> Float[torch.Tensor, "b c1 c2"]:
    return x[:, :, None] * y[:, None, :]

@jaxtyped(typechecker=beartype)
def scale_image(img: Float[torch.Tensor, "c h w"],
                scale: Float[torch.Tensor, ""]) -> Float[torch.Tensor, "c h w"]:
    return img * scale

@jaxtyped(typechecker=beartype)
def broadcast_add(a: Float[torch.Tensor, "#batch dim"],
                  b: Float[torch.Tensor, "#batch dim"]) -> Float[torch.Tensor, "#batch dim"]:
    return a + b

@jaxtyped(typechecker=beartype)
def flexible_net(x: Float[torch.Tensor, "*batch features"]) -> Float[torch.Tensor, "*batch features"]:
    return x * 2

x = torch.ones(3, 4)
y = torch.ones(3, 5)
print(batch_outer(x, y).shape)   # (3, 4, 5)

# Axis mismatch raises immediately:
# batch_outer(torch.ones(3, 4), torch.ones(7, 5))  # TypeCheckError: b=3 != b=7
```

--------------------------------

### Annotating jax.random.split output

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Demonstrates how to annotate the output of jax.random.split using the PRNGKeyArray type and a shape specification. This highlights the composability of Jaxtyping annotations.

```python
Shaped[PRNGKeyArray, "2"]
```

--------------------------------

### Creating Custom dtypes

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/advanced-features.md

Learn how to define your own data types using Jaxtyping's AbstractDtype and make_numpy_struct_dtype.

```APIDOC
## Creating your own dtypes

::: jaxtyping.AbstractDtype
    options:
        members: []

::: jaxtyping.make_numpy_struct_dtype
```

--------------------------------

### Using Local Variables in Symbolic Expressions

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Illustrates how to use local variables within symbolic expressions by enclosing them in curly braces. This enables dynamic shape annotations based on function arguments.

```python
def full(size: int, fill: float) -> Float[jax.Array, "{size}"]:
    return jax.numpy.full((size,), fill)
```

--------------------------------

### Load Jaxtyping IPython Extension and Set Typechecker

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/runtime-type-checking.md

Enable the Jaxtyping import hook in IPython environments and specify the runtime type checker to use.

```python
import jaxtyping
%load_ext jaxtyping
%jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
```

--------------------------------

### Symbolic Axis Expressions with JAX

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Shows how to use symbolic arithmetic for axis sizes in return type annotations, evaluated like f-strings. Supports subtraction, multiplication with variables, and dynamic sizing.

```python
import jax
import jax.numpy as jnp
from jaxtyping import Float, jaxtyped
from beartype import beartype

@jaxtyped(typechecker=beartype)
def remove_last(x: Float[jax.Array, "dim"]) -> Float[jax.Array, "dim-1"]:
    return x[:-1]

@jaxtyped(typechecker=beartype)
def tile(x: Float[jax.Array, "n"], times: int) -> Float[jax.Array, "{times}*n"]:
    return jnp.tile(x, times)

@jaxtyped(typechecker=beartype)
def full(size: int, fill: float) -> Float[jax.Array, "{size}"]:
    return jnp.full((size,), fill)

print(remove_last(jnp.ones(5)).shape)   # (4,)
print(tile(jnp.ones(3), 4).shape)       # (12,)
print(full(7, 3.14).shape)              # (7,)
```

--------------------------------

### Runtime Configuration of Jaxtyping

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Configure jaxtyping's runtime behavior using `jaxtyping.config`. Switches include disabling all checking and stripping typechecker stack frames from errors. Requires `jaxtyping`, `jax.numpy`, and `beartype`.

```python
from jaxtyping import config
import jax.numpy as jnp
from jaxtyping import Float, jaxtyped
from beartype import beartype

@jaxtyped(typechecker=beartype)
def add(x: Float[jnp.ndarray, "n"], y: Float[jnp.ndarray, "n"]) -> Float[jnp.ndarray, "n"]:
    return x + y

# Disable checking (e.g. in production / benchmarking)
config.update("jaxtyping_disable", True)
result = add(jnp.ones(3), jnp.ones(5))  # no error raised
print(result.shape)  # (5,) — mismatch was silently ignored

# Re-enable
config.update("jaxtyping_disable", False)
# add(jnp.ones(3), jnp.ones(5))  # now raises TypeCheckError again

# Cleaner error messages (hide typechecker's own frames)
config.update("jaxtyping_remove_typechecker_stack", True)
```

--------------------------------

### Configure Pytest for Jaxtyping Packages

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/runtime-type-checking.md

Specify packages to apply the Jaxtyping import hook using command-line arguments or configuration files.

```bash
pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
```

```toml
[tool.pytest.ini_options]
addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
```

```ini
[pytest]
addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
```

--------------------------------

### Duck-Type Array Support with Custom Dtypes

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Illustrates how to use jaxtyping with custom array-like objects that conform to duck-typing protocols (shape and dtype attributes). Shows defining a custom AbstractDtype and using it with jaxtyped.

```python
from jaxtyping import AbstractDtype, jaxtyped
from beartype import beartype

class MyDtype(AbstractDtype):
    dtypes = ["my_dtype"]

class MyDuckArray:
    @property
    def shape(self) -> tuple[int, ...]:
        return (3, 4)

    @property
    def dtype(self) -> str:
        return "my_dtype"

    def __mul__(self, other):
        return self

@jaxtyped(typechecker=beartype)
def process(arr: MyDtype[MyDuckArray, "3 4"]) -> MyDtype[MyDuckArray, "3 4"]:
    return arr * 2

obj = MyDuckArray()
print(isinstance(obj, MyDtype[MyDuckArray, "3 4"]))
print(process(obj))

# Mismatched shape fails:
class WrongShape(MyDuckArray):
    @property
    def shape(self):
        return (5, 4)

print(isinstance(WrongShape(), MyDtype[MyDuckArray, "3 4"])))
```

--------------------------------

### Allowed PyTree structure with matching leaf sizes

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/pytree.md

Demonstrates a valid usage of the `f` function where corresponding PyTree leaves have matching sizes, satisfying the path-dependent shape constraint.

```python
x0 = jnp.arange(3)
x1 = jnp.arange(5)

y0 = jnp.arange(3) + 1
y1 = jnp.arange(5) + 1

f((x0, x1), (y0, y1))  # x0 matches y0, and x1 matches y1. All good!
```

--------------------------------

### Applying `jaxtyped` Automatically with `install_import_hook`

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Use `install_import_hook` to automatically apply `@jaxtyped` to functions and classes within specified modules at import time. This is recommended for project-wide type checking without modifying individual files. It can be used as a context manager and supports manual uninstallation.

```python
# entry_point.py  ── apply to an entire package
from jaxtyping import install_import_hook

with install_import_hook("myproject", "beartype.beartype"):
    import myproject          # all functions in myproject.*
    import myproject.models   # are automatically @jaxtyped'd

# Alternative: manual uninstall
hook = install_import_hook(["myproject", "myutils"], "typeguard.typechecked")
import myproject
import myutils
hook.uninstall()

# ── Library __init__.py pattern ──────────────────────────────────────────────
# Inside mylib/__init__.py:
from jaxtyping import install_import_hook
with install_import_hook("mylib", "beartype.beartype"):
    from .model import MyModel
    from .train import train_step

# ── Pytest plugin (pyproject.toml) ───────────────────────────────────────────
# [tool.pytest.ini_options]
# addopts = "--jaxtyping-packages=myproject,beartype.beartype"
#
# CLI equivalent:
# pytest --jaxtyping-packages=myproject,beartype.beartype

# ── IPython / Jupyter ────────────────────────────────────────────────────────
import jaxtyping
# %load_ext jaxtyping
# %jaxtyping.typechecker beartype.beartype
```

--------------------------------

### Introspection

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/advanced-features.md

Discover how to introspect Jaxtyping types, checking for AbstractDtype, AbstractArray, and PyTree instances.

```APIDOC
## Introspection

::: jaxtyping.AbstractArray
    options:
        members: []

!!! info

    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)`., For example, `issubclass(Float32[jax.Array, "some shape"], AbstractArray)` will pass.

    You can check for pytrees by doing `issubclass(x, PyTree)`. For example, `issubclass(PyTree[int], PyTree)` will pass.
```

--------------------------------

### Debugging with `print_bindings`

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

A debugging helper to print the current values of bound axis names and PyTree structure names within a `jaxtyped`-decorated function. Requires `jax.numpy`, `jaxtyping`, and `beartype`.

```python
import jax.numpy as jnp
from jaxtyping import Float, jaxtyped, print_bindings
from beartype import beartype

@jaxtyped(typechecker=beartype)
def debug_shapes(x: Float[jnp.ndarray, "batch dim"],
                 y: Float[jnp.ndarray, "batch out"]) -> Float[jnp.ndarray, "batch out"]:
    print_bindings()
    # Prints something like:
    # The current values for each jaxtyping axis annotation are as follows.
    # batch=4
    # dim=8
    # out=3
    return y

debug_shapes(jnp.ones((4, 8)), jnp.ones((4, 3)))
```

--------------------------------

### Array Dtype Annotations with JAX

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Demonstrates using Float, Float32, Int, Bool, Shaped, Num, and jaxtyped for JAX arrays. Ensures specific dtypes and shapes are used, with type checking performed by beartype.

```python
import jax
import jax.numpy as jnp
from jaxtyping import Float, Float32, Int, Bool, Shaped, Num, jaxtyped
from beartype import beartype

# Accepts any floating-point JAX array of shape (batch, features)
@jaxtyped(typechecker=beartype)
def linear(w: Float[jax.Array, "out_dim in_dim"],
           x: Float[jax.Array, "batch in_dim"]
           ) -> Float[jax.Array, "batch out_dim"]:
    return x @ w.T

# Only 32-bit float
@jaxtyped(typechecker=beartype)
def relu(x: Float32[jax.Array, "batch dim"]) -> Float32[jax.Array, "batch dim"]:
    return jnp.maximum(x, 0)

# Bool array
@jaxtyped(typechecker=beartype)
def logical_not(mask: Bool[jax.Array, "n"]) -> Bool[jax.Array, "n"]:
    return ~mask

w = jnp.ones((4, 8))
x = jnp.ones((2, 8))
print(linear(w, x).shape)   # (2, 4)

bad = jnp.ones((2, 8), dtype=jnp.int32)
# linear(w, bad)  # raises TypeCheckError: expected Float, got int32
```

--------------------------------

### Run Tests with pytest

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/CONTRIBUTING.md

Execute all tests using pytest after making code changes to ensure functionality.

```bash
uv run pytest
```

--------------------------------

### Define Image and BatchImage types

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Defines type aliases for image and batched image arrays using Jaxtyping's Float and Shaped annotations. These demonstrate how to combine existing type aliases with new shape specifications.

```python
Image = Float[jax.Array, "channels height width"]
BatchImage = Float[Image, "batch"]
```

--------------------------------

### Using Class Attributes in Symbolic Expressions

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Shows how class attributes can be used within symbolic expressions for return type annotations, allowing for shapes to be defined by instance or class properties.

```python
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)
```

--------------------------------

### PyTree Annotation for Leaf Type and Structure Matching

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Annotates PyTree structures with leaf type constraints and structure names. Use `?` for path-dependent shape checking. Requires `jax`, `jax.numpy`, `jaxtyping`, and `beartype`.

```python
import jax
import jax.numpy as jnp
from jaxtyping import PyTree, Float, Shaped, jaxtyped
from beartype import beartype

# Basic PyTree leaf-type check
@jaxtyped(typechecker=beartype)
def sum_leaves(tree: PyTree[Float[jax.Array, "..."]]) -> float:
    return sum(x.sum() for x in jax.tree_util.tree_leaves(tree))

# Structure matching: x and y must have the same PyTree structure
@jaxtyped(typechecker=beartype)
def tree_add(x: PyTree[Float[jax.Array, "n"], "T"],
             y: PyTree[Float[jax.Array, "n"], "T"]
             ) -> PyTree[Float[jax.Array, "n"], "T"]:
    return jax.tree_util.tree_map(lambda a, b: a + b, x, y)

# Path-dependent shapes: corresponding leaves in x and y must match
@jaxtyped(typechecker=beartype)
def paired_scale(x: PyTree[Shaped[jax.Array, "?foo"], "T"],
                 y: PyTree[Shaped[jax.Array, "?foo"], "T"]):
    return jax.tree_util.tree_map(lambda a, b: a * b, x, y)

result = sum_leaves({"a": jnp.ones(3), "b": jnp.ones((2, 4))})
print(result)   # 11.0

x_tree = (jnp.ones(3), jnp.ones(5))
y_tree = (jnp.ones(3) * 2, jnp.ones(5) * 3)
print(tree_add(x_tree, y_tree))  # (array([3.]), array([4.]))

# tree_add(x_tree, {"a": jnp.ones(3)})  # TypeCheckError: structure mismatch
```

--------------------------------

### Path-dependent Shapes

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/pytree.md

Demonstrates the use of the '?' prefix for path-dependent axis sizes in PyTree annotations. This allows for specifying that corresponding leaves in different PyTrees must have the same size.

```APIDOC
## Path-dependent Shapes

The prefix `?` may be used to indicate that the axis size can depend on which leaf of a PyTree the array is at. For example:
```python
def f(
    x: PyTree[Shaped[jax.Array, "?foo"], "T"],
    y: PyTree[Shaped[jax.Array, "?foo"], "T"],
):
    pass
```
The above demands that `x` and `y` have matching PyTree structures (due to the `T` annotation), and that their leaves must all be one-dimensional arrays, *and that the corresponding pairs of leaves in `x` and `y` must have the same size as each other*.

Thus the following is allowed:
```python
x0 = jnp.arange(3)
x1 = jnp.arange(5)

y0 = jnp.arange(3) + 1
y1 = jnp.arange(5) + 1

f((x0, x1), (y0, y1))  # x0 matches y0, and x1 matches y1. All good!
```

But this is not:
```python
f((x1, x1), (y0, y1))  # x1 does not have a size matching y0!
```

Internally, all that is happening is that `foo` is replaced with `0foo` for the first leaf, `1foo` for the next leaf, etc., so that each leaf gets a unique version of the name.
```

--------------------------------

### Disallowed PyTree structure with non-matching leaf sizes

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/pytree.md

Illustrates an invalid usage of the `f` function where corresponding PyTree leaves have mismatched sizes, violating the path-dependent shape constraint.

```python
f((x1, x1), (y0, y1))  # x1 does not have a size matching y0!
```

--------------------------------

### Define function with path-dependent PyTree shapes

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/pytree.md

Use the '?' prefix to indicate that an axis size can depend on the leaf's position within a PyTree. This ensures corresponding leaves in different PyTrees have matching sizes.

```python
def f(
    x: PyTree[Shaped[jax.Array, "?foo"], "T"],
    y: PyTree[Shaped[jax.Array, "?foo"], "T"],
):
    pass
```

--------------------------------

### Define and Use Structured NumPy Dtype with Jaxtyping

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Defines a custom NumPy dtype for particles and uses it with a jaxtyped function for calculations. Ensure `make_numpy_struct_dtype` is imported.

```python
import numpy as np
from jaxtyping import make_numpy_struct_dtype, jaxtyped
from beartype import beartype

# Define a structured dtype
particle_dtype = np.dtype([("x", np.float32), ("y", np.float32), ("mass", np.float64)])
Particle = make_numpy_struct_dtype(particle_dtype)

@jaxtyped(typechecker=beartype)
def center_of_mass(particles: Particle[np.ndarray, "n"]) -> np.ndarray:
    total_mass = particles["mass"].sum()
    cx = (particles["x"] * particles["mass"]).sum() / total_mass
    cy = (particles["y"] * particles["mass"]).sum() / total_mass
    return np.array([cx, cy])

data = np.array(
    [(1.0, 2.0, 1.5), (3.0, 4.0, 0.5)],
    dtype=particle_dtype
)
print(center_of_mass(data))  # array([1.5, 2.5])
```

--------------------------------

### JAX Array Type Aliases

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Provides common type aliases for JAX arrays and array-like objects. These simplify type annotations by offering shorter, more descriptive names.

```python
jaxtyping.Array: alias for jax.Array
jaxtyping.ArrayLike: alias for jax.typing.ArrayLike
```

--------------------------------

### Using the `jaxtyped` Decorator for Function and Dataclass Type Checking

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

The `jaxtyped` decorator enforces type annotations on function arguments and dataclass fields at runtime. It supports various type checkers like `beartype` and can be used with `typechecker=None` for manual `isinstance` checks. It also supports context-manager style usage for temporary type-checking scopes.

```python
import jax
import jax.numpy as jnp
from dataclasses import dataclass
from jaxtyping import Float, jaxtyped
from beartype import beartype

# ── Functions ──────────────────────────────────────────────────────────────
@jaxtyped(typechecker=beartype)
def matrix_multiply(a: Float[jax.Array, "m k"],
                    b: Float[jax.Array, "k n"]
                    ) -> Float[jax.Array, "m n"]:
    return a @ b

# ── Dataclasses ─────────────────────────────────────────────────────────────
@jaxtyped(typechecker=beartype)
@dataclass
class Batch:
    images: Float[jax.Array, "b c h w"]
    labels: Float[jax.Array, "b"]

# ── Manual isinstance checks (typechecker=None) ──────────────────────────────
@jaxtyped(typechecker=None)
def check_manually(x):
    assert isinstance(x, Float[jax.Array, "batch dim"]), "bad shape or dtype"
    return x

# ── Context-manager style ────────────────────────────────────────────────────
def at_global_scope(x):
    with jaxtyped("context"):
        assert isinstance(x, Float[jax.Array, "n"])

a = jnp.ones((3, 4))
b = jnp.ones((4, 5))
print(matrix_multiply(a, b).shape)   # (3, 5)

batch = Batch(images=jnp.ones((8, 3, 32, 32)), labels=jnp.ones(8))

# matrix_multiply(jnp.ones((3, 4)), jnp.ones((9, 5)))
# → TypeCheckError: k=4 does not equal k=9
```

--------------------------------

### JAX PRNGKeyArray Type Alias

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Provides a type alias for PRNGKeyArray, covering both new-style typed keys and old-style keys. This is useful for annotating functions that handle JAX random key splitting or generation.

```python
jaxtyping.PRNGKeyArray: alias for jaxtyping.Key[jax.Array, ""] | jaxtyping.UInt32[jax.Array, "2"]
```

--------------------------------

### Creating NumPy Structured Dtype Classes with `make_numpy_struct_dtype`

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

The `make_numpy_struct_dtype` function generates a jaxtyping dtype class for NumPy structured arrays. This enables type annotations for arrays with named fields, similar to named tuples.

```python
import numpy as np
from jaxtyping import make_numpy_struct_dtype, jaxtyped
from beartype import beartype
```

--------------------------------

### Define Typed Matrix Multiplication Function

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/index.md

Use Jaxtyping annotations to define a function that accepts and returns typed tensors with specific shapes and dtypes. This ensures type safety at runtime when used with compatible type-checking packages.

```python
from jaxtyping import Float
from torch import Tensor

# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Tensor, "dim1 dim2"],
                    y: Float[Tensor, "dim2 dim3"]
                  ) -> Float[Tensor, "dim1 dim3"]:
    ...
```

--------------------------------

### JAX PRNGKey Type

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Defines the type for JAX PRNG keys. This annotation can be used for functions that produce or consume JAX random keys.

```python
jaxtyping.Key, which is the dtype of jax.random.key's. For example jax.random.key(...) produces a jaxtyping.Key[jax.Array, ""]
```

--------------------------------

### Control Array Name Format in Error Messages

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Use `set_array_name_format` and `get_array_name_format` to control how array type names are rendered in error messages. The default is 'dtype_and_shape'.

```python
from jaxtyping import set_array_name_format, get_array_name_format

print(get_array_name_format())  # "dtype_and_shape"

set_array_name_format("array")  # show only the array class in errors
# ... run your code ...
set_array_name_format("dtype_and_shape")  # restore default
```

--------------------------------

### Defining Custom Dtype Categories with `AbstractDtype`

Source: https://context7.com/patrick-kidger/jaxtyping/llms.txt

Subclass `AbstractDtype` to create custom dtype categories for use in type annotations. The `dtypes` attribute specifies the accepted dtype strings. This allows for more flexible type checking, especially when dealing with different precision levels or custom array types.

```python
import numpy as np
from jaxtyping import AbstractDtype, jaxtyped
from beartype import beartype

class Float16or32(AbstractDtype):
    """Accepts float16 or float32 arrays."""
    dtypes = ["float16", "float32"]

class MyCustomDtype(AbstractDtype):
    """Matches a custom duck-type array dtype string."""
    dtypes = ["my_special_dtype"]

@jaxtyped(typechecker=beartype)
def low_precision_add(x: Float16or32[np.ndarray, "n"],
                      y: Float16or32[np.ndarray, "n"]
                      ) -> Float16or32[np.ndarray, "n"]:
    return x + y

a = np.ones(5, dtype=np.float32)
b = np.ones(5, dtype=np.float16)
# low_precision_add(a, b) works if they have the same dtype at call-site;
# b/c jaxtyping checks each argument independently for dtype membership.

# Introspection
from jaxtyping import AbstractArray, Float32
import jax
print(issubclass(Float32, AbstractDtype))                       # True
print(issubclass(Float32[jax.Array, "a b"], AbstractArray))    # True
```

--------------------------------

### JAX Scalar Type Aliases

Source: https://github.com/patrick-kidger/jaxtyping/blob/main/docs/api/array.md

Defines type aliases for scalar JAX arrays and scalar array-like objects. These use the Shaped annotation with an empty shape string to denote scalars.

```python
jaxtyping.Scalar: alias for jaxtyping.Shaped[jax.Array, ""]
jaxtyping.ScalarLike: alias for jaxtyping.Shaped[jax.typing.ArrayLike, ""]
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.