Add support for torch.export models

[tolleybot]

TorchSharp currently lacks support for loading and executing PyTorch models exported via torch.export and compiled with AOTInductor (.pt2 files). While TorchSharp supports TorchScript (JIT) models via torch.jit.load(), it cannot load the newer AOTInductor-compiled format introduced in PyTorch 2.x. This is becoming critical as the PyTorch ecosystem transitions from ONNX to torch.export for model deployment.

torch.export is PyTorch 2.x's modern, recommended approach for model export, using TorchDynamo for symbolic tracing and producing optimized single-graph representations with only ATen-level operations. To load these models from C++, an additional AOTInductor compilation step is required.

• Link for PyTorch implementation:

  • torch.export Documentation
  • torch.export Tutorial
  • torch.export API Reference
  • PT2 Archive Spec
  • AOTInductor Documentation

• Example code for PyTorch implementation:

import torch
import torch.nn as nn

# Define a simple model
class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# Export and compile the model for AOTInductor
model = MyModule()
example_inputs = (torch.randn(1, 10),)
exported_program = torch.export.export(model, example_inputs)

# Compile and package for C++ loading (this is the key step)
torch._inductor.aoti_compile_and_package(
    exported_program, package_path="model.pt2"
)

# The resulting .pt2 can be loaded in C++ via AOTIModelPackageLoader
# Note: torch.export.save() produces a different .pt2 format
# that is Python-only and cannot be loaded from C++

Important distinction: There are two different .pt2 formats:

• torch.export.save() produces a Python-only .pt2 file that cannot be loaded from C++
• torch._inductor.aoti_compile_and_package() produces an AOTInductor .pt2 file that can be loaded from C++ via AOTIModelPackageLoader

This feature targets the AOTInductor format only.

Describe the solution you'd like

Add support for loading and executing AOTInductor-compiled .pt2 models in TorchSharp. The solution should include:

1. Loading API - Ability to load AOTInductor-compiled .pt2 files:

// New namespace for torch.export
using var exportedProgram = torch.export.load("model.pt2");

2. Execution API - Ability to run forward pass (inference only):

var inputTensor = torch.randn(1, 10);
var result = exportedProgram.run(inputTensor);

3. Implementation approach:
   • Use PyTorch C++ API's torch::inductor::AOTIModelPackageLoader (requires LibTorch 2.9+)
   • Leverage existing TorchSharp infrastructure (three-layer architecture: C# → PInvoke → C++ wrapper → LibTorch)
   • Inference-only API (no training, parameter updates, or device movement)
   • Models must be compiled with torch._inductor.aoti_compile_and_package() in Python

Technical considerations:

• LibTorch 2.9+ includes torch::inductor::AOTIModelPackageLoader for loading AOTInductor-compiled models
• .pt2 files from torch.export.save() are Python-only and cannot be loaded in C++
• Only .pt2 files from torch._inductor.aoti_compile_and_package() work in C++
• Models are compiled for specific device (CPU/CUDA) at build time
• Can follow similar implementation pattern as ScriptModule (see src/TorchSharp/JIT/ScriptModule.cs)

Environment Information:

• OS: All platforms (Windows, Linux, MacOS)
• Package Type: All packages (torchsharp-cpu, torchsharp-cuda-windows, torchsharp-cuda-linux)
• Version: Latest (currently 0.105.x) and future versions

Additional context

Why this matters:

• torch.export is the recommended path forward for PyTorch model deployment (PyTorch 2.x+)
• Without this support, .NET developers cannot consume the latest PyTorch model formats
• TorchScript (JIT) is becoming legacy; torch.export is the future

Related Issues:

• Support for non Tensor types for ScriptModule #782 - Support for non Tensor types for ScriptModule
• Support more types for ScriptModule #771 - Support more types for ScriptModule (closed)
• Load TorchScript modules #644 - Load TorchScript modules (merged) - Similar pattern can be followed
• Will TorchSharp support exporting ONNX models in the future? #1281 - Will TorchSharp support exporting ONNX models in the future? (closed)

Comparison: TorchScript vs torch.export:

Feature	TorchScript (JIT)	torch.export (AOTInductor)
File Format	.pt	.pt2
PyTorch Era	1.x	2.x+ (current)
Status in TorchSharp	✅ Supported	❌ Not supported
Python Export Step	torch.jit.save()	torch.export.export() → aoti_compile_and_package()
Use Case	Research	Production deployment

[alinpahontu2912]

Hey @tolleybot thanks for the PR and the help, would you like to give it a try again, since the new release is based on libtorch 2.10?

[tolleybot]

Definitely