Metadata-Version: 2.4
Name: coreai-torch
Version: 0.4.0
Summary: Convert PyTorch models to CoreAI format
License: Copyright 2026 Apple Inc.
        
        Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
        
        1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
        
        2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
        
        3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
        
        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
        
Project-URL: Homepage, https://github.com/apple/coreai-torch
Project-URL: Repository, https://github.com/apple/coreai-torch
Project-URL: Issues, https://github.com/apple/coreai-torch/issues
Classifier: License :: OSI Approved :: BSD License
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: coreai-core==1.0.0b1
Requires-Dist: ml-dtypes
Requires-Dist: networkx
Requires-Dist: numpy
Requires-Dist: packaging
Requires-Dist: scipy
Requires-Dist: sympy
Requires-Dist: torch<=2.11.0,>=2.8.0
Requires-Dist: typing-extensions
Requires-Dist: strenum
Requires-Dist: rich<16.0,>=13.0
Provides-Extra: test
Requires-Dist: coremltools; extra == "test"
Requires-Dist: filecheck; extra == "test"
Requires-Dist: mlx; sys_platform == "darwin" and extra == "test"
Requires-Dist: mlx_lm; sys_platform == "darwin" and extra == "test"
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-asyncio; extra == "test"
Requires-Dist: pytest-randomly; extra == "test"
Requires-Dist: pytest-rerunfailures; extra == "test"
Requires-Dist: pytest-sugar; extra == "test"
Requires-Dist: pytest-xdist; extra == "test"
Requires-Dist: torchaudio; extra == "test"
Requires-Dist: torchvision; extra == "test"
Requires-Dist: transformers==4.57.3; extra == "test"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: shibuya==2026.5.19; extra == "docs"
Requires-Dist: myst-nb>=1.0; extra == "docs"
Requires-Dist: nbmake>=1.5.0; extra == "docs"
Requires-Dist: ghp-import>=2.1.0; extra == "docs"
Requires-Dist: nest-asyncio>=1.5.0; extra == "docs"
Provides-Extra: dev
Requires-Dist: deptry>=0.20; extra == "dev"
Requires-Dist: nbstripout; extra == "dev"
Requires-Dist: pre-commit>=4.0; extra == "dev"
Requires-Dist: ruff>=0.12.0; extra == "dev"
Dynamic: license-file

# Core AI PyTorch Extensions (coreai-torch)

Core AI PyTorch Extensions (`coreai-torch`) is a Python package that bridges PyTorch and Core AI. Use it to bring up an existing PyTorch model into Core AI IR, or to author Core AI models directly from PyTorch by composing the built-in composite op library (`coreai_torch.composite_ops`), authoring new ops via `register_torch_lowering`, and authoring inline Metal GPU kernels via `TorchMetalKernel`. The resulting IR can be compiled and executed efficiently by the Core AI inference stack.

🔗 Jump to: [Getting started](#getting-started) · [Documentation](#documentation) · [Contributing](#contributing) · [Support](#support) · [License](#license)

## Overview

`coreai-torch` traverses a `torch.export.ExportedProgram` and produces Core AI
IR — the same IR consumed by the Core AI compiler and runtime. The public entry
point is `TorchConverter`, which lowers PyTorch operators to Core AI dialect
operations, preserves location and module-stack information for debugging, and
provides extension points for custom Metal kernels and submodule
externalization.

## Getting started

### Installation

```bash
pip install coreai-torch
```

Or from source with [uv](https://docs.astral.sh/uv/):

```bash
uv sync
```

### Usage

`TorchConverter` accepts models in two forms. Pick based on what you have and whether you need externalization:

| You have | Use |
|---|---|
| A decomposed `ExportedProgram` | `add_exported_program()` |
| An `nn.Module` + need externalization | `add_pytorch_module()` with `externalize_modules` |
| An `nn.Module`, no externalization | Either method — `add_exported_program()` is more direct |

#### From an ExportedProgram

You export and decompose the model yourself, then pass the `ExportedProgram` directly.
Use `get_decomp_table()` so that composite ops (`instance_norm`, `pixel_shuffle`, `scaled_dot_product_attention`) are preserved for optimal runtime performance.

```python
import torch
from coreai_torch import TorchConverter, get_decomp_table

model = ...  # your nn.Module
model.eval()

# Export and decompose — this is your responsibility
ep = torch.export.export(model, args=(torch.randn(1, 3, 224, 224),))
ep = ep.run_decompositions(get_decomp_table())

# Convert to Core AI IR
converter = TorchConverter().add_exported_program(ep)
coreai_program = converter.to_coreai()
coreai_program.optimize()
```

#### From an nn.Module

Pass your model and an `export_fn` that returns a decomposed `ExportedProgram`. This is equivalent to calling `add_exported_program()` with the result of `export_fn`.

```python
import coreai_torch
from coreai_torch import TorchConverter

model = ...  # your nn.Module
model.eval()
sample = (torch.randn(1, 3, 224, 224),)

converter = TorchConverter().add_pytorch_module(
    model,
    export_fn=lambda m: torch.export.export(m, args=sample).run_decompositions(
        coreai_torch.get_decomp_table()
    ),
)
coreai_program = converter.to_coreai()
coreai_program.optimize()
```

## Documentation

- API reference and conversion guides — see the [`docs/`](docs/) directory
- Supported model types and ops — see the operator coverage notebooks under [`docs/`](docs/)

### Building docs locally

```bash
uv sync --extra docs
uv run jupyter-book build docs/
open docs/_build/html/index.html
```

## Development

### Running tests

```bash
uv sync --extra test
uv run pytest tests/ -n auto
```

### Testing notebooks

```bash
uv run pytest docs/ --nbmake -v
```

## Contributing

We welcome contributions within a defined scope. Please read
[CONTRIBUTING.md](CONTRIBUTING.md) carefully before opening a pull request or
issue — particularly the section on contribution scope.

## Support

- [GitHub Issues](../../issues) — Bug reports, feature requests, and questions

## Security and code of conduct

Security vulnerability reporting and the Code of Conduct for this project are
governed at the org level via the
[Apple Open Source `.github` repository](https://github.com/apple/.github).

## License

This project is licensed under the [BSD 3-Clause License](LICENSE).

## Related projects

- [Core AI](https://developer.apple.com/documentation/coreai) — Apple's on-device AI inference stack
- [Core AI Optimization](https://github.com/apple/coreai-optimization) — model compression for deployment on Apple Silicon
- [Core AI Models](https://github.com/apple/coreai-models) — ready-to-run optimized models, Python reproduction scripts, and Swift utilities for on-device integration
