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 · Documentation · Contributing · Support · 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

pip install coreai-torch

Or from source with uv:

uv sync

Usage

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

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.

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.

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/ directory
• Supported model types and ops — see the operator coverage notebooks under docs/

Building docs locally

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

Development

Setting up pre-commit hooks

This repo uses pre-commit to run linting and formatting checks automatically. Install the hooks once after cloning:

pre-commit install --hook-type pre-commit --hook-type pre-push

Running tests

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

Testing notebooks

uv run pytest docs/ --nbmake -v

Contributing

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

Support

• GitHub 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.

License

This project is licensed under the BSD 3-Clause License.

Related projects

• Core AI — Apple's on-device AI inference stack
• Core AI Optimization — model compression for deployment on Apple Silicon
• Core AI Models — ready-to-run optimized models, Python reproduction scripts, and Swift utilities for on-device integration