ExTorch.Export (extorch v0.3.0)

Copy Markdown

Read and introspect PyTorch ExportedProgram .pt2 archives.

This module provides a pure-Elixir reader for .pt2 files produced by torch.export.save(). It can extract the model graph, weight metadata, and raw weight tensors without requiring Python or C++ ExportedProgram support.

Python export workflow

import torch

model = MyModel()
model.eval()
exported = torch.export.export(model, (example_input,))
torch.export.save(exported, "model.pt2")

Elixir usage

# Load and run inference directly
model = ExTorch.Export.load("model.pt2")
output = ExTorch.Export.forward(model, [input])

# Or read schema and weights separately
schema = ExTorch.Export.read_schema("model.pt2")
weights = ExTorch.Export.read_weights("model.pt2")

# Generate DSL source code
IO.puts(ExTorch.Export.to_elixir("model.pt2", "MyModel"))

Note

This reads .pt2 files from torch.export.save, NOT from aoti_compile_and_package. AOTI-compiled .pt2 files don't contain the graph or separable weights -- use ExTorch.AOTI for those.

Summary

Functions

forward(model, inputs)

Run inference on a loaded Export model.

load(path)

Load an exported .pt2 model for inference.

read_schema(path)

Read the model schema from an exported .pt2 archive.

read_weights(path)

Load weight tensors from an exported .pt2 archive.

to_elixir(path, module_name \\ "MyModel")

Generate an ExTorch.NN.Module DSL definition from an exported .pt2 archive.

Functions

forward(model, inputs)

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

Run inference on a loaded Export model.

Interprets the ATen computation graph, dispatching each operation to the corresponding ExTorch tensor function.

Args

  • model (ExTorch.Export.Model) - the loaded model.
  • inputs ([ExTorch.Tensor]) - input tensors, matching the model's user inputs.

Returns

The output tensor (or list of tensors for multi-output models).

Example

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

load(path)

@spec load(String.t()) :: ExTorch.Export.Model.t()

Load an exported .pt2 model for inference.

Reads the graph and weights, and prepares the model for forward/2.

Args

  • path (String) - path to the .pt2 file from torch.export.save.

Returns

An %ExTorch.Export.Model{} struct.

Example

model = ExTorch.Export.load("model.pt2")
output = ExTorch.Export.forward(model, [input_tensor])

read_schema(path)

@spec read_schema(String.t()) :: map()

Read the model schema from an exported .pt2 archive.

Returns a map with:

  • :graph - the computation graph as a list of node maps
  • :inputs - graph input names
  • :outputs - graph output names
  • :weights - weight metadata (name → shape, dtype, requires_grad)

read_weights(path)

@spec read_weights(String.t()) :: %{required(String.t()) => ExTorch.Tensor.t()}

Load weight tensors from an exported .pt2 archive.

Returns a map of %{fqn => %ExTorch.Tensor{}}.

to_elixir(path, module_name \\ "MyModel")

@spec to_elixir(String.t(), String.t()) :: String.t()

Generate an ExTorch.NN.Module DSL definition from an exported .pt2 archive.

Maps ATen operations in the graph to ExTorch NN layer types where possible.

Args

  • path - path to the .pt2 file.
  • module_name - name for the generated Elixir module.