xref: /aosp_15_r20/external/executorch/extension/android_test/add_model.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1import torch
2from executorch.exir import to_edge
3from torch.export import export
4
5
6# Start with a PyTorch model that adds two input tensors (matrices)
7class Add(torch.nn.Module):
8    def __init__(self):
9        super(Add, self).__init__()
10
11    def forward(self, x: torch.Tensor, y: torch.Tensor):
12        return x + y
13
14
15# 1. torch.export: Defines the program with the ATen operator set.
16aten_dialect = export(Add(), (torch.ones(1), torch.ones(1)))
17
18# 2. to_edge: Make optimizations for Edge devices
19edge_program = to_edge(aten_dialect)
20
21# 3. to_executorch: Convert the graph to an ExecuTorch program
22executorch_program = edge_program.to_executorch()
23
24# 4. Save the compiled .pte program
25with open("add.pte", "wb") as file:
26    file.write(executorch_program.buffer)
27