jax-shapeguard 0.3.0__tar.gz

jax_shapeguard-0.3.0/.github/workflows/lint.yml

@@ -0,0 +1,37 @@
+name: Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Set up Python
+        run: uv python install 3.11
+
+      - name: Install dependencies
+        run: uv sync --dev
+
+      - name: Install linting tools
+        run: uv pip install ruff mypy
+
+      - name: Run ruff check
+        run: uv run ruff check shapeguard tests
+
+      - name: Run ruff format check
+        run: uv run ruff format --check shapeguard tests
+
+      - name: Run mypy
+        run: uv run mypy shapeguard --ignore-missing-imports

jax_shapeguard-0.3.0/.github/workflows/publish.yml

@@ -0,0 +1,29 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write  # Required for trusted publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Set up Python
+        run: uv python install 3.11
+
+      - name: Build package
+        run: uv build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

jax_shapeguard-0.3.0/.github/workflows/test.yml

@@ -0,0 +1,39 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --dev
+
+      - name: Run tests
+        run: uv run pytest -v --cov=shapeguard --cov-report=xml
+
+      - name: Upload coverage
+        uses: codecov/codecov-action@v4
+        if: matrix.python-version == '3.11'
+        with:
+          files: coverage.xml
+          fail_ci_if_error: false

jax_shapeguard-0.3.0/.gitignore

@@ -0,0 +1,44 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+.nox/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db

jax_shapeguard-0.3.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jayendra Parmar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jax_shapeguard-0.3.0/MILESTONES.md

@@ -0,0 +1,241 @@
+# ShapeGuard Milestones
+
+## Design Decisions
+
+- **Dim identity**: Same object required (`n = Dim("n")` must be reused)
+- **Performance**: Dev-only tool, prioritize error quality over speed
+- **Backend priority**: NumPy-first, JAX support in v0.2
+
+---
+
+## Milestone 1: Core Foundation (v0.1-alpha)
+
+### Goal
+Minimal working library with symbolic dimensions, shape checking, and decorator.
+
+### Files
+```
+shapeguard/
+  __init__.py       # Public API exports
+  core.py           # Dim class, UnificationContext
+  spec.py           # Shape specification matching
+  decorator.py      # @expects decorator
+  errors.py         # ShapeGuardError
+  _compat.py        # Backend detection
+tests/
+  test_core.py
+  test_spec.py
+  test_decorator.py
+  test_errors.py
+  conftest.py
+pyproject.toml
+```
+
+### Deliverables
+- [x] `Dim` class with identity-based unification
+- [x] `UnificationContext` for tracking bindings across arguments
+- [x] `ShapeSpec` matching: concrete `(3, 4)`, symbolic `(n, m)`, wildcard `(None, 4)`
+- [x] `@expects` decorator for input validation
+- [x] `ShapeGuardError` with function, argument, expected, actual, reason
+- [x] `check_shape(x, spec, name)` standalone function
+- [x] Backend-agnostic shape extraction (works with any `.shape` attribute)
+- [x] Unit tests with 90%+ coverage (91% achieved)
+
+### API Surface
+```python
+from shapeguard import Dim, expects, check_shape, ShapeGuardError
+
+n, m = Dim("n"), Dim("m")
+
+@expects(x=(n, m), y=(m,))
+def forward(x, y):
+    return x @ y
+
+check_shape(arr, (n, 128), name="input")
+```
+
+---
+
+## Milestone 2: ML-Practical Features (v0.1-beta)
+
+### Goal
+Ergonomic features for real ML workflows.
+
+### Deliverables
+- [x] `Batch` dimension (always first, flexible size per call)
+- [x] Ellipsis support `(..., n, m)` for variable leading dims
+- [x] `ShapeContext` manager for grouped checks with shared bindings
+- [x] Improved error messages with binding trace (92% coverage)
+
+### API Additions
+```python
+from shapeguard import Batch, ShapeContext
+
+B = Batch()
+
+@expects(x=(B, n, m))
+def layer(x): ...
+
+@expects(x=(..., n, m))
+def normalize(x): ...
+
+with ShapeContext() as ctx:
+    ctx.check(x, (n, m), "x")
+    ctx.check(y, (m, k), "y")
+```
+
+---
+
+## Milestone 3: JAX Integration (v0.2)
+
+### Goal
+Seamless JAX compatibility including JIT behavior.
+
+### Deliverables
+- [x] JIT/tracing detection
+- [x] Configurable JIT modes: `skip`, `warn`, `check`
+- [x] PyTree shape specs for nested params
+- [ ] Performance benchmarks (deferred)
+
+### API Additions
+```python
+from shapeguard import expects, config
+
+config.jit_mode = "skip"  # Global setting
+
+@expects(x=(B, n, m), jit_mode="static")  # Per-function
+@jax.jit
+def forward(x): ...
+
+@expects(
+    params={"weights": (n, m), "bias": (m,)},
+    x=(B, n)
+)
+def apply(params, x): ...
+```
+
+---
+
+## Milestone 4: Broadcasting Support (v0.2)
+
+### Goal
+Explicit broadcasting inspection and validation.
+
+### Deliverables
+- [x] `broadcast_shape()` for concrete shapes
+- [x] `explain_broadcast()` step-by-step explainer
+- [ ] `_broadcast=True` option in `@expects` (deferred)
+
+### API Additions
+```python
+from shapeguard import broadcast_shape, explain_broadcast
+
+broadcast_shape((3, 1), (1, 4))  # → (3, 4)
+broadcast_shape(a, b)            # From arrays
+
+explain_broadcast((3, 1, 4), (5, 4))
+# Broadcasting (3, 1, 4) with (5, 4):
+#   Dim 0: 3 (from left only)
+#   Dim 1: 1 → 5 (broadcast)
+#   Dim 2: 4 = 4 (match)
+#   Result: (3, 5, 4)
+```
+
+---
+
+## Milestone 5: Output Contracts (v0.3)
+
+### Goal
+Validate function outputs, not just inputs.
+
+### Deliverables
+- [x] `@ensures` decorator for output validation
+- [x] `@contract` combined decorator
+- [x] Tuple/dict output support
+
+### API Additions
+```python
+from shapeguard import expects, ensures, contract
+
+@expects(a=(n, m), b=(m, k))
+@ensures(result=(n, k))
+def matmul(a, b):
+    return a @ b
+
+@contract(
+    inputs={"a": (n, m), "b": (m, k)},
+    output=(n, k)
+)
+def matmul(a, b):
+    return a @ b
+```
+
+---
+
+## Milestone 6: ML Helpers (v0.3)
+
+### Goal
+Domain-specific helpers for common ML patterns.
+
+### Deliverables
+- [ ] Pre-defined dims: `B`, `T`, `C`, `H`, `W`, `D`
+- [ ] `attention_shapes()` helper
+- [ ] `conv_output_shape()` calculator
+
+### API Additions
+```python
+from shapeguard.ml import B, T, C, H, W, D
+from shapeguard.ml import attention_shapes, conv_output_shape
+
+@expects(x=(B, T, D))
+def transformer_layer(x): ...
+
+@expects(**attention_shapes(B, heads, seq_q, seq_k, d_k))
+def attention(q, k, v): ...
+
+out_shape = conv_output_shape(
+    input=(B, C, 224, 224),
+    kernel=(3, 3),
+    stride=2,
+    padding=1
+)
+```
+
+---
+
+## Milestone 7: Testing Utilities (v0.4)
+
+### Goal
+Property-based testing support.
+
+### Deliverables
+- [ ] Hypothesis strategies for shaped arrays
+- [ ] `verify_contract()` auto-test generator
+- [ ] pytest plugin
+
+### API Additions
+```python
+from shapeguard.testing import arrays, verify_contract
+import hypothesis
+
+@hypothesis.given(x=arrays(shape=(n, m), n=(1, 100), m=(1, 100)))
+def test_normalize(x):
+    result = normalize(x)
+    assert result.shape == x.shape
+
+verify_contract(matmul, samples=100)
+```
+
+---
+
+## Summary Timeline
+
+| Milestone | Version | Status |
+|-----------|---------|--------|
+| 1. Core Foundation | v0.1-alpha | ✅ Complete (91% coverage) |
+| 2. ML Features | v0.1-beta | ✅ Complete (92% coverage) |
+| 3. JAX Integration | v0.2 | ✅ Complete (92% coverage) |
+| 4. Broadcasting | v0.2 | ✅ Complete |
+| 5. Output Contracts | v0.3 | ✅ Complete (91% coverage) |
+| 6. ML Helpers | v0.3 | 🔲 Not started |
+| 7. Testing Utils | v0.4 | 🔲 Not started |

jax_shapeguard-0.3.0/PKG-INFO

@@ -0,0 +1,69 @@
+Metadata-Version: 2.4
+Name: jax-shapeguard
+Version: 0.3.0
+Summary: Runtime shape contracts and diagnostics for NumPy and JAX
+Project-URL: Homepage, https://github.com/jayendra13/jax-shape-guard
+Project-URL: Repository, https://github.com/jayendra13/jax-shape-guard
+Author-email: Jayendra Parmar <jayendra0parmar@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: debugging,jax,ml,numpy,shapes,validation
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Provides-Extra: jax
+Requires-Dist: jax>=0.4; extra == 'jax'
+Requires-Dist: jaxlib>=0.4; extra == 'jax'
+Description-Content-Type: text/markdown
+
+# ShapeGuard
+
+[![Tests](https://github.com/jayendra13/jax-shape-guard/actions/workflows/test.yml/badge.svg)](https://github.com/jayendra13/jax-shape-guard/actions/workflows/test.yml)
+[![Lint](https://github.com/jayendra13/jax-shape-guard/actions/workflows/lint.yml/badge.svg)](https://github.com/jayendra13/jax-shape-guard/actions/workflows/lint.yml)
+[![PyPI version](https://img.shields.io/pypi/v/shapeguard.svg)](https://pypi.org/project/shapeguard/)
+[![Python versions](https://img.shields.io/pypi/pyversions/shapeguard.svg)](https://pypi.org/project/shapeguard/)
+[![License](https://img.shields.io/github/license/jayendra13/jax-shape-guard.svg)](https://github.com/jayendra13/jax-shape-guard/blob/main/LICENSE)
+
+Runtime shape contracts and diagnostics for NumPy and JAX.
+
+## Installation
+
+```bash
+pip install shapeguard
+```
+
+## Quick Start
+
+```python
+from shapeguard import Dim, expects
+
+n, m, k = Dim("n"), Dim("m"), Dim("k")
+
+@expects(a=(n, m), b=(m, k))
+def matmul(a, b):
+    return a @ b
+```
+
+When shapes don't match, you get clear errors:
+
+```
+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]
+```
+
+## License
+
+MIT

jax_shapeguard-0.3.0/README.md

@@ -0,0 +1,42 @@
+# ShapeGuard
+
+[![Tests](https://github.com/jayendra13/jax-shape-guard/actions/workflows/test.yml/badge.svg)](https://github.com/jayendra13/jax-shape-guard/actions/workflows/test.yml)
+[![Lint](https://github.com/jayendra13/jax-shape-guard/actions/workflows/lint.yml/badge.svg)](https://github.com/jayendra13/jax-shape-guard/actions/workflows/lint.yml)
+[![PyPI version](https://img.shields.io/pypi/v/shapeguard.svg)](https://pypi.org/project/shapeguard/)
+[![Python versions](https://img.shields.io/pypi/pyversions/shapeguard.svg)](https://pypi.org/project/shapeguard/)
+[![License](https://img.shields.io/github/license/jayendra13/jax-shape-guard.svg)](https://github.com/jayendra13/jax-shape-guard/blob/main/LICENSE)
+
+Runtime shape contracts and diagnostics for NumPy and JAX.
+
+## Installation
+
+```bash
+pip install shapeguard
+```
+
+## Quick Start
+
+```python
+from shapeguard import Dim, expects
+
+n, m, k = Dim("n"), Dim("m"), Dim("k")
+
+@expects(a=(n, m), b=(m, k))
+def matmul(a, b):
+    return a @ b
+```
+
+When shapes don't match, you get clear errors:
+
+```
+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]
+```
+
+## License
+
+MIT

jax_shapeguard-0.3.0/pyproject.toml

@@ -0,0 +1,94 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["shapeguard"]
+
+[project]
+name = "jax-shapeguard"
+version = "0.3.0"
+description = "Runtime shape contracts and diagnostics for NumPy and JAX"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+authors = [
+    { name = "Jayendra Parmar", email = "jayendra0parmar@gmail.com" }
+]
+keywords = ["numpy", "jax", "shapes", "validation", "debugging", "ml"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Typing :: Typed",
+]
+
+[project.optional-dependencies]
+jax = [
+    "jax>=0.4",
+    "jaxlib>=0.4",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=7.0",
+    "pytest-cov>=4.0",
+    "numpy>=1.20",
+    "jax>=0.4",
+    "jaxlib>=0.4",
+    "ruff>=0.4",
+    "mypy>=1.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/jayendra13/jax-shape-guard"
+Repository = "https://github.com/jayendra13/jax-shape-guard"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-v --tb=short"
+
+[tool.coverage.run]
+source = ["shapeguard"]
+branch = true
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+    "raise NotImplementedError",
+]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100
+
+[tool.ruff.lint]
+select = [
+    "E",      # pycodestyle errors
+    "W",      # pycodestyle warnings
+    "F",      # pyflakes
+    "I",      # isort
+    "B",      # flake8-bugbear
+    "UP",     # pyupgrade
+]
+ignore = [
+    "E501",   # line too long (handled by formatter)
+]
+
+[tool.ruff.lint.isort]
+known-first-party = ["shapeguard"]
+
+[tool.mypy]
+python_version = "3.10"
+warn_return_any = true
+warn_unused_ignores = true
+disallow_untyped_defs = false
+ignore_missing_imports = true

jax_shapeguard-0.3.0/shapeguard/__init__.py

@@ -0,0 +1,61 @@
+"""
+ShapeGuard: Runtime shape contracts and diagnostics for NumPy and JAX.
+
+Basic usage:
+    from shapeguard import Dim, expects, check_shape
+
+    n, m = Dim("n"), Dim("m")
+
+    @expects(x=(n, m), y=(m,))
+    def forward(x, y):
+        return x @ y
+
+ML workflows:
+    from shapeguard import Batch, ShapeContext
+
+    B = Batch()
+
+    @expects(x=(B, n, m))
+    def layer(x): ...
+
+    # Ellipsis for variable leading dims
+    @expects(x=(..., n, m))
+    def normalize(x): ...
+
+    # Grouped checks
+    with ShapeContext() as ctx:
+        ctx.check(x, (n, m), "x")
+        ctx.check(y, (m, k), "y")
+"""
+
+from shapeguard.broadcast import broadcast_shape, explain_broadcast
+from shapeguard.config import config
+from shapeguard.context import ShapeContext
+from shapeguard.core import Batch, Dim, UnificationContext
+from shapeguard.decorator import contract, ensures, expects
+from shapeguard.errors import BroadcastError, OutputShapeError, ShapeGuardError
+from shapeguard.spec import check_shape
+
+__version__ = "0.3.0"
+
+__all__ = [
+    # Core
+    "Dim",
+    "Batch",
+    "UnificationContext",
+    # Validation
+    "expects",
+    "ensures",
+    "contract",
+    "check_shape",
+    "ShapeContext",
+    # Broadcasting
+    "broadcast_shape",
+    "explain_broadcast",
+    # Configuration
+    "config",
+    # Errors
+    "ShapeGuardError",
+    "OutputShapeError",
+    "BroadcastError",
+]