Metadata-Version: 2.4
Name: jaxtyc
Version: 0.7.2
Summary: Static array shape checking for JAX powered by eval_shape
Project-URL: Homepage, https://github.com/BeeGass/jaxtyc
Project-URL: Documentation, https://beegass.github.io/jaxtyc/
Project-URL: Repository, https://github.com/BeeGass/jaxtyc
Project-URL: Changelog, https://github.com/BeeGass/jaxtyc/blob/main/CHANGELOG.md
Project-URL: Bug Tracker, https://github.com/BeeGass/jaxtyc/issues
Author-email: Bryan Gass <44324535+BeeGass@users.noreply.github.com>
License-Expression: MIT
License-File: LICENSE
Keywords: jax,jaxtyping,lsp,machine-learning,shape-checking,static-analysis,type-checking
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Quality Assurance
Classifier: Typing :: Typed
Requires-Python: !=3.12,<3.15,>=3.11
Requires-Dist: jaxtyping<0.4,>=0.2.28
Requires-Dist: pygls<3,>=2.0
Provides-Extra: all
Requires-Dist: einops<0.9,>=0.8; extra == 'all'
Requires-Dist: equinox<0.14,>=0.11; extra == 'all'
Requires-Dist: flax<0.10.3,>=0.10.0; extra == 'all'
Requires-Dist: watchfiles<2,>=1.0; extra == 'all'
Provides-Extra: cpu
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'cpu'
Requires-Dist: jax[cpu]>=0.9.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'cpu'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'cpu'
Provides-Extra: cpu-nightly
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'cpu-nightly'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'cpu-nightly'
Provides-Extra: cuda12
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'cuda12'
Requires-Dist: jax[cuda12]>=0.9.0; (sys_platform == 'linux') and extra == 'cuda12'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'cuda12'
Provides-Extra: cuda12-nightly
Requires-Dist: jax-cuda12-pjrt; (sys_platform == 'linux') and extra == 'cuda12-nightly'
Requires-Dist: jax-cuda12-plugin[with-cuda]; (sys_platform == 'linux') and extra == 'cuda12-nightly'
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'cuda12-nightly'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'cuda12-nightly'
Provides-Extra: cuda13
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'cuda13'
Requires-Dist: jax[cuda13]>=0.9.0; (sys_platform == 'linux') and extra == 'cuda13'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'cuda13'
Provides-Extra: cuda13-nightly
Requires-Dist: jax-cuda13-pjrt; (sys_platform == 'linux') and extra == 'cuda13-nightly'
Requires-Dist: jax-cuda13-plugin[with-cuda]; (sys_platform == 'linux') and extra == 'cuda13-nightly'
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'cuda13-nightly'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'cuda13-nightly'
Provides-Extra: einops
Requires-Dist: einops<0.9,>=0.8; extra == 'einops'
Provides-Extra: equinox
Requires-Dist: equinox<0.14,>=0.11; extra == 'equinox'
Provides-Extra: flax
Requires-Dist: flax<0.10.3,>=0.10.0; extra == 'flax'
Provides-Extra: mac
Requires-Dist: jax-mps>=0.9.12; (sys_platform == 'darwin' and platform_machine == 'arm64' and python_version == '3.13') and extra == 'mac'
Requires-Dist: jax>=0.4.34; (sys_platform == 'darwin') and extra == 'mac'
Requires-Dist: jaxlib>=0.4.34; (sys_platform == 'darwin') and extra == 'mac'
Provides-Extra: rocm
Requires-Dist: jax-rocm7-pjrt; (sys_platform == 'linux' and platform_machine == 'x86_64' and python_version < '3.14') and extra == 'rocm'
Requires-Dist: jax-rocm7-plugin; (sys_platform == 'linux' and platform_machine == 'x86_64' and python_version < '3.14') and extra == 'rocm'
Requires-Dist: jax>=0.7.1; (sys_platform == 'linux' and platform_machine == 'x86_64' and python_version < '3.14') and extra == 'rocm'
Requires-Dist: jaxlib>=0.8.0; (sys_platform == 'linux' and platform_machine == 'x86_64' and python_version < '3.14') and extra == 'rocm'
Provides-Extra: tpu
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'tpu'
Requires-Dist: jax[tpu]>=0.9.0; (sys_platform == 'linux') and extra == 'tpu'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'tpu'
Provides-Extra: tpu-nightly
Requires-Dist: jax>=0.9.0; (sys_platform == 'linux') and extra == 'tpu-nightly'
Requires-Dist: jax[tpu]>=0.9.0; (sys_platform == 'linux') and extra == 'tpu-nightly'
Requires-Dist: jaxlib>=0.9.0; (sys_platform == 'linux') and extra == 'tpu-nightly'
Provides-Extra: watch
Requires-Dist: watchfiles<2,>=1.0; extra == 'watch'
Description-Content-Type: text/markdown

[![PyPI](https://img.shields.io/pypi/v/jaxtyc)](https://pypi.org/project/jaxtyc/) [![Python](https://img.shields.io/pypi/pyversions/jaxtyc)](https://pypi.org/project/jaxtyc/) [![CI](https://github.com/BeeGass/jaxtyc/actions/workflows/ci.yml/badge.svg)](https://github.com/BeeGass/jaxtyc/actions/workflows/ci.yml) [![License](https://img.shields.io/github/license/BeeGass/jaxtyc)](https://github.com/BeeGass/jaxtyc/blob/main/LICENSE)

# jaxtyc

Static array shape checking for JAX powered by `jax.eval_shape`.

Reads [jaxtyping](https://docs.kidger.site/jaxtyping/) annotations and verifies shapes at analysis time -- no runtime cost, no FLOPs. Each named dimension is assigned a unique prime number, making shape mismatches unambiguous.

<p align="center">
  <img src="docs/assets/vscode-inlay-hints.png" alt="VS Code inlay hints showing sharding annotations and shape overlays" width="600">
</p>

<p align="center">
  <img src="docs/assets/cli-diagnostics.png" alt="CLI diagnostics showing shape mismatches in Claude Code" width="600">
</p>

<p align="center">
  <a href="docs/assets/demo.mov">Watch the demo video</a>
</p>

## Features

- **Zero runtime cost** -- `jax.eval_shape` only; no arrays allocated, no computation executed
- **Prime-based symbolic shapes** -- each dimension name maps to a unique prime (>= 101), so `d_in != d_out` is guaranteed
- **10 diagnostic rules** -- shape/rank mismatch, cross-function propagation, parameter consistency, tuple return checking, trace errors
- **Inline suppressions** -- `# jaxtyc: ignore` and `# jaxtyc: ignore[rule-name]`
- **LSP server** -- diagnostics, hover, CodeLens, go-to-definition, references, rename, code actions, completion, semantic tokens, inlay hints, signature help, linked editing, folding, call hierarchy
- **LSP multiplexer** -- `jaxtyc mux` runs ty/pyright + jaxtyc behind a single stdio pipe
- **CLI with 4 output formats** -- `full`, `concise`, `json`, `github` (inline PR annotations)
- **Flax NNX + Equinox support** -- traces bound methods on module instances
- **Configurable via `pyproject.toml`** -- severity threshold, rule ignoring, file exclusion, einops preferences

## Installation

```bash
pip install jaxtyc
# or
uv add jaxtyc
```

**Extras:**

| Extra | Installs | Use case |
|-------|----------|----------|
| `jaxtyc[watch]` | `watchfiles` | `jaxtyc watch` -- re-check on file save |
| `jaxtyc[flax]` | `flax >=0.10` | Flax NNX module tracing |
| `jaxtyc[equinox]` | `equinox >=0.11` | Equinox module tracing |
| `jaxtyc[einops]` | `einops >=0.8` | einops-style fix suggestions + inlay hints with pattern dim names |
| `jaxtyc[all]` | All of the above | Everything |

## Quick Start

```python
# model.py
import jax.numpy as jnp
from jaxtyping import Array, Float

def linear(
    x: Float[Array, "batch seq d_in"],
    w: Float[Array, "d_in d_out"],
) -> Float[Array, "batch seq d_out"]:
    return jnp.matmul(x, w.T)  # Bug: .T swaps dims, produces (batch, seq, d_in)
```

```bash
$ jaxtyc check model.py
model.py:8:0: error[shape-mismatch]
  Shape mismatch in return of `linear`
    Expected: (batch, seq, d_out)
    Got:      (batch, seq, d_in)

Found 1 error(s) in 1 function(s) checked (0.03s)
```

## Editor Integration

### VS Code

Install the [jaxtyc extension](editors/vscode/):

```bash
cd editors/vscode && npm install && npm run bundle
npx @vscode/vsce package --allow-missing-repository
code --install-extension jaxtyc-*.vsix
```

Or use the justfile: `just vscode-update`

The extension auto-discovers your Python environment (`.venv`, `VIRTUAL_ENV`, or `jaxtyc` on PATH) and starts the LSP server automatically. Supports multi-root workspaces with per-folder LSP clients. Includes jaxtyping snippets, a trace visualization webview, and a status bar quick pick menu.

### Other Editors

jaxtyc works in any editor that supports LSP (Neovim, Helix, etc.). See the [editor setup docs](docs/editors/editors.md) for configuration.

## CLI

```
jaxtyc check <paths>...          # Shape-check files or directories
jaxtyc trace <file.py::func>     # Trace intermediate shapes through a function
jaxtyc watch <paths>...          # Watch and re-check on change
jaxtyc lsp                       # Start the LSP server (stdio)
jaxtyc mux                       # Start the LSP multiplexer (ty/pyright + jaxtyc)
jaxtyc version                   # Print version
```

## Documentation

Full docs at [beegass.github.io/jaxtyc](https://beegass.github.io/jaxtyc/).

## Contributing

Contributions are welcome! See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines.

## License

MIT
