Metadata-Version: 2.5
Name: jax-pyrefly-stubs
Version: 1.4.0.dev3
Summary: JAX type stubs with array-shape tracking, maintained by Pyrefly
Project-URL: homepage, https://pyrefly.org
Project-URL: documentation, https://pyrefly.org/en/docs/
License: MIT License
        
        Copyright (c) Meta Platforms, Inc. and affiliates.
        
        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.
License-File: LICENSE
Keywords: jax,stubs,typechecker,typechecking
Classifier: Development Status :: 4 - Beta
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Typing :: Stubs Only
Requires-Python: >=3.12
Requires-Dist: pyrefly-shape-extensions==1.4.0.dev3
Description-Content-Type: text/markdown

# Pyrefly JAX shape stubs

This package is a PEP 561 stub-only distribution. It installs the `jax-stubs`
stub package so Pyrefly can discover shape-aware stubs for the runtime `jax`
package without replacing or shadowing JAX itself. It is versioned in lockstep
with Pyrefly and depends on the matching `pyrefly-shape-extensions` package.

Shape-typed fixture stubs for a subset of JAX. This is a starting point for
working with the JAX core team rather than a complete model: it covers array
creation, broadcasting arithmetic, `matmul`, `reshape`, `transpose`, the
`axis`/`keepdims` reductions, and the elementwise activations in `jax.nn`.
Dimensions are modeled and dtypes are not, so adding dtypes later means modeling
JAX's own defaults (`float32` and `int32`) rather than copying the NumPy stubs.

Shape rules use the type-level DSL, `@type_shape_dsl_function`, exclusively. A
rule the DSL cannot express yet returns a gradual shape instead. Where a rule is
imprecise for that reason, the stub says so at the definition, along with
whether a fix is expected.

`TENSOR_SHAPES_CONTRIBUTING.md` at the repository root covers the workflow, and
`tensor-shapes/run_tests.py` runs the tests. Pyrefly checks the stubs themselves
and every `test/test_*.py`, and those same test files then run against real
JAX, so a stub that is self-consistent but wrong still fails.

Before submitting a change, format this package with Ruff from the repository
root:

```bash
uv tool run --from ruff==0.16.5 ruff format tensor-shapes/jax-pyrefly-stubs
```

Anything not listed above is simply absent rather than modeled loosely, so it is
reported as a missing attribute rather than inferred gradually.
