ExTorch.AOTI (extorch v0.3.0)

Copy Markdown

Load and run inference on AOTInductor-compiled models (.pt2 packages).

AOTI (Ahead-of-Time Inductor) models are compiled through torch.export and torch._inductor.aoti_compile_and_package(), producing optimized .pt2 packages with fused kernels. These models offer better inference throughput than TorchScript but do not support introspection or weight extraction.

Python export workflow

import torch
from torch._inductor import aoti_compile_and_package

model = MyModel()
model.eval()
example_input = torch.randn(1, 10)

exported = torch.export.export(model, (example_input,))
aoti_compile_and_package(exported, package_path="model.pt2")

Elixir inference

model = ExTorch.AOTI.load("model.pt2")
input = ExTorch.randn({1, 10})
[output] = ExTorch.AOTI.forward(model, [input])

Availability

AOTI support requires a libtorch build that includes the inductor runtime. Check with ExTorch.AOTI.available?/0.

Summary

Functions

available?()

Check if AOTI support is available in the current libtorch build.

constant_names(model)

Get the fully-qualified names of constants (parameters/buffers) in the model.

forward(model, inputs)

Run inference on an AOTI model.

load(path, opts \\ [])

Load an AOTI-compiled model from a .pt2 package.

metadata(model)

Get metadata from an AOTI model as a map.

Functions

available?()

@spec available?() :: boolean()

Check if AOTI support is available in the current libtorch build.

constant_names(model)

@spec constant_names(ExTorch.AOTI.Model.t()) :: [String.t()]

Get the fully-qualified names of constants (parameters/buffers) in the model.

Returns

A list of strings like ["fc1.weight", "fc1.bias", ...].

forward(model, inputs)

@spec forward(ExTorch.AOTI.Model.t(), [ExTorch.Tensor.t()]) :: [ExTorch.Tensor.t()]

Run inference on an AOTI model.

Args

  • model (ExTorch.AOTI.Model) - the loaded model.
  • inputs ([ExTorch.Tensor]) - list of input tensors.

Returns

A list of output tensors.

load(path, opts \\ [])

@spec load(
  String.t(),
  keyword()
) :: ExTorch.AOTI.Model.t()

Load an AOTI-compiled model from a .pt2 package.

Args

  • path (String) - path to the .pt2 file.
  • opts (keyword) - optional arguments:
    • :model_name (String) - name of the model within the package. Default: "model".
    • :device_index (integer) - device index for CUDA. Default: -1 (CPU).

Returns

An %ExTorch.AOTI.Model{} struct.

metadata(model)

@spec metadata(ExTorch.AOTI.Model.t()) :: %{required(String.t()) => String.t()}

Get metadata from an AOTI model as a map.

Returns

A map of %{String.t() => String.t()} metadata key-value pairs.