I built ShapeGuard because shape errors in numerical code are uniquely painful. They're silent (wrong shapes often produce garbage instead of crashing), late (errors surface deep in XLA,
not where the bug is), and cryptic (shapes (3,4) and (5,3) not aligned — but why should they match?).
ShapeGuard lets you declare shape contracts on functions using symbolic dimensions:
from shapeguard import Dim, expects, ensures
n, m, k = Dim("n"), Dim("m"), Dim("k")
@expects(a=(n, m), b=(m, k))
@ensures(result=(n, k))
def matmul(a, b):
return a @ b
When shapes don't match, the error traces bindings back to their source:
ShapeGuardError:
function: matmul
argument: b
expected: (m, k)
actual: (5, 7)
reason: dimension 'm' bound to 4 from a.shape[1], but got 5 from b.shape[0]
bindings: {n=3 (from a[0]), m=4 (from a[1])}
The key idea is unification — the same Dim object used across arguments must resolve to the same integer. ShapeGuard tracks where each binding came from, so conflicts pinpoint the exact
source.
What it does:
- @expects / @ensures / @contract — input and output shape validation
- Symbolic Dim with cross-argument unification
- Batch() dims and ... ellipsis for ML patterns
- broadcast_shape() and explain_broadcast() for debugging broadcasting
- Configurable JIT modes (check / warn / skip) for JAX
- ML helpers: pre-defined dims (B, T, C, H, W, D), attention_shapes(), conv_output_shape()
What it doesn't do:
- No dtype checking (jaxtyping does this well)
- No named-tensor wrapper (Haliax's approach)
- Not a replacement for static type checking
It's zero-dependency, drop-in (works with existing code — just add decorators), and the motivation came from analyzing 40 real JAX GitHub issues where users hit cryptic shape errors.
PyPI: pip install jax-shapeguard
Would love feedback on the API design, error message format, and whether this would actually help your workflow. What shape debugging pain points am I missing?