bearshape

Runtime shape and dtype checking for NumPy, JAX, PyTorch, CuPy, and tree-structured containers, powered by beartype.

from beartype import beartype
from bearshape import N, C
from bearshape.numpy import F32


@beartype
def normalize(x: F32[N, C]) -> F32[N, C]:
  return x / x.sum(axis=1, keepdims=True)

bearshape turns annotations such as F32[N, C], F32Like[~B, C], F32[Value("size")], and Tree[F32[N], T] into runtime-validated contracts. Named dimensions are shared within a function call, so mismatched shapes fail at the boundary instead of later in array code.

Install

pip install bearshape

bearshape keeps the root import lightweight. Install the array backend packages you use explicitly:

pip install bearshape numpy
pip install bearshape numpy torch
pip install bearshape numpy jax
pip install bearshape numpy cupy
pip install bearshape numpy optree

What It Checks

• strict array type, dtype, and shape contracts
• backend-aware Like[...] conversion checks
• scalar-like values and constrained runtime Value(...) dimensions
• tree leaf and structure annotations through JAX or OpTree
• annotation syntax exercised by pyright, mypy, and ty fixtures

Public Surface

Import dimensions, Value, Scalar, DtypeSpec, check, and check_context from bearshape.

Import backend aliases from backend modules:

from bearshape.numpy import F32, F32Like
from bearshape.jax import Tree
from bearshape.torch import I64

The root package does not import NumPy or any backend. Backend modules require their own runtime dependencies.

Development

uv sync
uv run prek run -a
uv run pytest -n auto tests/
uv run pytest -n auto tests/test_typecheck.py
uv run pyright src tests/typing
uv run mypy src tests/typing
uv run ty check src tests/typing

CuPy runtime tests require a CUDA-capable environment and are deferred on CPU-only machines.

See the documentation site for the full user guide and API reference.