justinchuby

import torch
import onnx_ir as ir

class ControlFlowModel(torch.nn.Module):
    def forward(self, x):
        def times_2(x):
            return x * 2

        def neg(x):
            return -x

justinchuby

"""Export to ONNX.

transformers_version == "4.52.0"
"""

import onnx_diagnostic.tasks.text_generation
import torch
from transformers import AutoConfig, AutoModel

import onnxscript

justinchuby

import onnx_ir as ir
import onnx


def create_model():
    """Create a model that has a unsorted node with subgraph that uses a value defined later."""
    a = ir.Value(name="a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3)))
    b = ir.Value(name="b", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((3, 4)))
    b_out = ir.Value(name="b_out", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((3, 4)))
    c = ir.Value(name="c", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((4, 5)))

justinchuby

torch.export non strict mode uses "make_fx [which] uses __torch_dispatch__ to trace under the hood. this is where it creates the fx nodes. AOTAutograd also calls into make_fx, but before that it also does some things related to functionalization. Since export is now "Training IR", it no longer does functionalization, so we just directly call make_fx."

justinchuby

"""Verify the cast values"""
import os
import onnx
import onnx_ir as ir

DIR = "onnx/backend/test/data/node/"

def verify_one_case(path: str):
    test_name = os.path.basename(path)
    input_path = os.path.join(path, "test_data_set_0", "input_0.pb")

justinchuby

https://github.com/iree-org/iree-turbine/blob/main/iree/turbine/aot/fx_programs.py
Also ai-edge torch exporter

justinchuby

from ai_edge_torch.odml_torch.export import exported_program_to_mlir
import torch

class PowModel(torch.nn.Module):

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x ** 0.5

model = PowModel()
print(model(torch.tensor(2)))

justinchuby

import torch
from onnx_diagnostic import torch_export_patches
from onnxscript.ir.passes.common import clear_metadata_and_docstring

from transformers import AttentionInterface, AutoModelForCausalLM, AutoTokenizer
from transformers.cache_utils import DynamicCache


# Get position_ids from attention_mask
def get_position_ids(attention_mask: torch.Tensor, use_past_kv: bool):

justinchuby

# Owner(s): ["module: onnx"]
"""Unit LLM tests for the onnx dynamo exporter."""

from __future__ import annotations

from typing import Any
import logging

import transformers

justinchuby

import logging
import torch
from torch_geometric.nn import GAT

logger = logging.getLogger(__name__)
logging.getLogger('torch.onnx').setLevel(logging.INFO)

logger.info("Prepare model")
num_features = 23
num_classes = 12