Metadata-Version: 2.4
Name: pyrefly-jax-stubs
Version: 0.10.2.1
Summary: JAX type stubs with native Pyrefly tensor-shape tracking
Project-URL: Homepage, https://github.com/bdelwood/pyrefly-jax-stubs
Project-URL: Repository, https://github.com/bdelwood/pyrefly-jax-stubs
Project-URL: Issues, https://github.com/bdelwood/pyrefly-jax-stubs/issues
Author-email: Brodi Elwood <bdelwood@fas.harvard.edu>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Typing :: Stubs Only
Requires-Python: >=3.13
Requires-Dist: jax==0.10.2
Requires-Dist: pyrefly-shape-extensions==1.2.0.dev2
Description-Content-Type: text/markdown

# pyrefly-jax-stubs

Native Pyrefly shape types for JAX arrays.

This PEP 561 partial stub package preserves JAX's shipped typing and specializes array-bearing APIs with Pyrefly's `SizeTuple`, `SymVar`, and shape syntax:

```python
from jax import Array
from shape_extensions import SizeTuple, SymVar

def rows[N: SymVar](x: Array[[N, 3]]) -> Array[[N]]: ...
def identity[S: SizeTuple](x: Array[S]) -> Array[S]: ...
```

Pyrefly discovers the installed stubs automatically. APIs not included here continue to use the typing shipped with JAX.

## Installation

Install the stubs and Pyrefly as development dependencies:

```bash
uv add --dev pyrefly==1.2.0.dev2 pyrefly-jax-stubs==0.10.2.1
```

The package requires Python 3.13 or newer because its signatures use type-parameter defaults. Release `0.10.2.1` targets JAX 0.10.2 and is tested with Pyrefly 1.2.0.dev2.

## Development

A local checkout can be checked and built with:

```bash
uv sync
just lint
just test
just build
```

The copied and adapted JAX declarations retain their original copyright notices and remain under the Apache License 2.0.

## Versioning

Releases use `<jax-version>.<stub-revision>`. The package pins JAX exactly so an upstream API change cannot silently invalidate the overlay.
