justinchuby

import torch


class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(5, 10)

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

justinchuby

# Adapted from
# https://github.com/pytorch/pytorch/blob/b505e8647547f029d0f7df408ee5f2968f757f89/test/test_public_bindings.py#L523
# Original code PyTorch license https://github.com/pytorch/pytorch/blob/main/LICENSE
# Modifications Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from __future__ import annotations

import importlib
import itertools
import os

justinchuby

value_name = self.name if self.name is not None else "anonymous:" + str(id(self))

if producer is not None:
            producer_text = producer.name if producer.name is not None else "anonymous:" + str(id(producer))
        else:
            producer_text = "None"

justinchuby

"""Utilities for traversing the IR graph."""

from __future__ import annotations

__all__ = [
    "RecursiveGraphIterator",
]

from typing import Callable, Iterator, Reversible

justinchuby

ir_version: 8
producer_name: "pytorch"
producer_version: "2.3.0"
graph {
  node {
    output: "_val_2"
    name: "Constant_0"
    op_type: "Constant"
    attribute {
      name: "value_floats"

justinchuby

from onnxscript import ir
import onnx

model_proto = onnx.load("model.onnx")
# (not const) -> cast to 16 -> cast to 32 -> Op
model = ir.serde.deserialize_model(model_proto)


def is_cast(node: ir.Node, dtype: ir.DataType):
    if node.op_type != "Cast":

justinchuby

<
    ir_version=8,
    opset_imports={'pkg.onnxscript.torch_lib': 1, 'pkg.torch.2.4.0a0+gitd56ab7b': 1, 'pkg.transformers.4.37.2': 1, '': 18, 'pkg.onnxscript.torch_lib.common': 1},
    producer_name='pytorch',
    producer_version='2.4.0',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,

justinchuby

"""Convenience methods for constructing (and manipulating?) the IR."""

from __future__ import annotations

import collections.abc
from typing import Any, Mapping, Sequence

from onnxrewriter.experimental.intermediate_representation import _ir

justinchuby

# https://aiinfra.visualstudio.com/PublicPackages/_artifacts/feed/ORT-Nightly/PyPI/ort-nightly/overview
--index-url=https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
ort-nightly==1.17.0.dev20240118001

justinchuby

# onnxscript/tests/function_libs/torch_lib/dynamo_export_test.py

import copy
import inspect
import itertools
import sys
import unittest

import torch
from torch.onnx import ExportOptions