Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/_export/__init__.py

@@ -0,0 +1,189 @@
+# mypy: allow-untyped-defs
+import copy
+import dataclasses
+import functools
+import io
+import json
+import logging
+import os
+import re
+import sys
+import types
+import warnings
+import weakref
+import zipfile
+from collections import OrderedDict
+from contextlib import contextmanager
+from functools import lru_cache
+
+from typing import Any, Optional, TYPE_CHECKING, Union
+from collections.abc import Callable
+from unittest.mock import patch
+
+import torch
+import torch.fx
+import torch.utils._pytree as pytree
+
+from torch._dispatch.python import enable_python_dispatcher
+from torch._guards import compile_context
+from torch._utils_internal import log_export_usage
+from torch.export._tree_utils import reorder_kwargs
+from torch.export.graph_signature import (
+    ArgumentSpec,
+    ConstantArgument,
+    ExportGraphSignature,
+    InputKind,
+    InputSpec,
+    OutputKind,
+    OutputSpec,
+    SymIntArgument,
+    SymBoolArgument,
+    SymFloatArgument,
+    TensorArgument,
+)
+from torch.fx import traceback as fx_traceback
+from torch.fx._compatibility import compatibility
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
+
+from .wrappers import _wrap_submodules
+from .utils import _materialize_cpp_cia_ops
+from . import config
+
+if TYPE_CHECKING:
+    from torch._C._aoti import AOTIModelContainerRunner
+
+log = logging.getLogger(__name__)
+
+@dataclasses.dataclass
+class ExportDynamoConfig:
+    """
+    Manage Export-specific configurations of Dynamo.
+    """
+    allow_rnn: bool = True
+
+
+# We only want to print this once to avoid flooding logs in workflows where aot_compile_warning
+# is called multiple times.
+@lru_cache
+def aot_compile_warning():
+
+    log.warning("+============================+")
+    log.warning("|     !!!   WARNING   !!!    |")
+    log.warning("+============================+")
+    log.warning(
+        "torch._export.aot_compile()/torch._export.aot_load() is being deprecated, please switch to "
+        "directly calling torch._inductor.aoti_compile_and_package(torch.export.export())/"
+        "torch._inductor.aoti_load_package() instead.")
+
+
+def aot_compile(
+    f: Callable,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
+    *,
+    dynamic_shapes: dict[str, Any] | None = None,
+    options: dict[str, Any] | None = None,
+    remove_runtime_assertions: bool = False,
+    disable_constraint_solver: bool = False,
+    same_signature: bool = True,
+) -> list[Any] | str:
+    """
+    Note: this function is not stable yet
+
+    Traces either an nn.Module's forward function or just a callable with PyTorch
+    operations inside, generates executable cpp code from the program, and returns
+    the path to the generated shared library
+
+    Args:
+        f: the `nn.Module` or callable to trace.
+
+        args: example positional inputs.
+
+        kwargs: optional example keyword inputs.
+
+        dynamic_shapes: Should either be:
+            1) a dict from argument names of ``f`` to their dynamic shape specifications,
+            2) a tuple that specifies dynamic shape specifications for each input in original order.
+            If you are specifying dynamism on keyword args, you will need to pass them in the order that
+            is defined in the original function signature.
+
+            The dynamic shape of a tensor argument can be specified as either
+            (1) a dict from dynamic dimension indices to :func:`Dim` types, where it is
+            not required to include static dimension indices in this dict, but when they are,
+            they should be mapped to None; or (2) a tuple / list of :func:`Dim` types or None,
+            where the :func:`Dim` types correspond to dynamic dimensions, and static dimensions
+            are denoted by None. Arguments that are dicts or tuples / lists of tensors are
+            recursively specified by using mappings or sequences of contained specifications.
+
+        options: A dictionary of options to control inductor
+
+        disable_constraint_solver: Whether the dim constraint solver must be disabled.
+
+    Returns:
+        Path to the generated shared library
+    """
+    from torch.export._trace import _export_to_torch_ir
+    from torch._inductor.decomposition import select_decomp_table
+    from torch._inductor import config as inductor_config
+
+    aot_compile_warning()
+
+    if inductor_config.is_predispatch:
+        gm = torch.export._trace._export(f, args, kwargs, dynamic_shapes, pre_dispatch=True).module()
+    else:
+        # We want to export to Torch IR here to utilize the pre_grad passes in
+        # inductor, which run on Torch IR.
+        with torch._export.config.patch(use_new_tracer_experimental=True):
+            gm = _export_to_torch_ir(
+                f,
+                args,
+                kwargs,
+                dynamic_shapes,
+                disable_constraint_solver=disable_constraint_solver,
+                same_signature=same_signature,
+                # Disabling this flag, because instead we can rely on the mapping
+                # dynamo_flat_name_to_original_fqn which is coming from Dynamo.
+                restore_fqn=False,
+            )
+
+    with torch.no_grad():
+        so_path = torch._inductor.aot_compile(gm, args, kwargs, options=options)  # type: ignore[arg-type]
+
+    if not isinstance(so_path, (str, list)):
+        raise AssertionError(f"expected str or list, got {type(so_path)}")
+    return so_path
+
+def aot_load(so_path: str, device: str) -> Callable:
+    """
+    Loads a shared library generated by aot_compile and returns a callable
+
+    Args:
+        so_path: Path to the shared library
+
+    Returns:
+        A callable
+    """
+    aot_compile_warning()
+
+    if device == "cpu":
+        runner: AOTIModelContainerRunner = torch._C._aoti.AOTIModelContainerRunnerCpu(so_path, 1)
+    elif device == "cuda" or device.startswith("cuda:"):
+        runner = torch._C._aoti.AOTIModelContainerRunnerCuda(so_path, 1, device)
+    elif device == "xpu" or device.startswith("xpu:"):
+        runner = torch._C._aoti.AOTIModelContainerRunnerXpu(so_path, 1, device)
+    elif device == "mps" or device.startswith("mps:"):
+        runner = torch._C._aoti.AOTIModelContainerRunnerMps(so_path, 1)
+    else:
+        raise RuntimeError("Unsupported device " + device)
+
+    def optimized(*args, **kwargs):
+        call_spec = runner.get_call_spec()
+        in_spec = pytree.treespec_loads(call_spec[0])
+        out_spec = pytree.treespec_loads(call_spec[1])
+        flat_inputs = pytree.tree_flatten((args, reorder_kwargs(kwargs, in_spec)))[0]
+        flat_inputs = [x for x in flat_inputs if isinstance(x, torch.Tensor)]
+        flat_outputs = runner.run(flat_inputs)
+        return pytree.tree_unflatten(flat_outputs, out_spec)
+
+    return optimized

kronos-venv/lib/python3.12/site-packages/torch/_export/config.py

@@ -0,0 +1,45 @@
+"""
+Configuration module for torch.export.export.
+
+This module contains various configuration flags and settings that control torch.export's
+behavior, including:
+- Runtime behavior flags
+- Debugging and development options
+"""
+
+import sys
+from typing import Any, TYPE_CHECKING
+
+from torch._environment import is_fbcode
+from torch.utils._config_module import install_config_module
+
+
+# this flag controls whether we use new functional tracer. It
+# should be True in the long term.
+use_new_tracer_experimental = True
+
+# this flag is used to control whether we want to instrument
+# fake tensor creation to track potential leaks. It is off
+# by default, but user can turn it on to debug leaks.
+detect_non_strict_fake_tensor_leaks = False
+
+# error on potentially pre-dispatch/non-strict tracing limitation
+# this type of error usually happens when we encounter an op
+# that we don't know how to proxy, resulting in untracked fake tensors
+error_on_lifted_constant_tensors = True
+
+# enable auto_functionalized_v2 in export
+# We turn this off in fbcode due to downstream users not
+# being ready to handle auto_functionalized_v2.
+enable_auto_functionalized_v2_for_export = not is_fbcode()
+
+use_legacy_dynamo_graph_capture = True
+
+
+if TYPE_CHECKING:
+    from torch.utils._config_typing import *  # noqa: F401, F403
+
+    def _make_closure_patcher(**changes: Any) -> Any: ...
+
+
+install_config_module(sys.modules[__name__])

kronos-venv/lib/python3.12/site-packages/torch/_export/db/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

kronos-venv/lib/python3.12/site-packages/torch/_export/db/case.py

@@ -0,0 +1,177 @@
+# mypy: allow-untyped-defs
+import inspect
+import re
+import string
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+from types import ModuleType
+
+import torch
+
+_TAGS: dict[str, dict[str, Any]] = {
+    "torch": {
+        "cond": {},
+        "dynamic-shape": {},
+        "escape-hatch": {},
+        "map": {},
+        "dynamic-value": {},
+        "operator": {},
+        "mutation": {},
+    },
+    "python": {
+        "assert": {},
+        "builtin": {},
+        "closure": {},
+        "context-manager": {},
+        "control-flow": {},
+        "data-structure": {},
+        "standard-library": {},
+        "object-model": {},
+    },
+}
+
+
+class SupportLevel(Enum):
+    """
+    Indicates at what stage the feature
+    used in the example is handled in export.
+    """
+
+    SUPPORTED = 1
+    NOT_SUPPORTED_YET = 0
+
+
+ArgsType = tuple[Any, ...]
+
+
+def check_inputs_type(args, kwargs):
+    if not isinstance(args, tuple):
+        raise ValueError(
+            f"Expecting args type to be a tuple, got: {type(args)}"
+        )
+    if not isinstance(kwargs, dict):
+        raise ValueError(
+            f"Expecting kwargs type to be a dict, got: {type(kwargs)}"
+        )
+    for key in kwargs:
+        if not isinstance(key, str):
+            raise ValueError(
+                f"Expecting kwargs keys to be a string, got: {type(key)}"
+            )
+
+def _validate_tag(tag: str):
+    parts = tag.split(".")
+    t = _TAGS
+    for part in parts:
+        if not set(part) <= set(string.ascii_lowercase + "-"):
+            raise AssertionError(f"Tag contains invalid characters: {part}")
+        if part in t:
+            t = t[part]
+        else:
+            raise ValueError(f"Tag {tag} is not found in registered tags.")
+
+
+@dataclass(frozen=True)
+class ExportCase:
+    example_args: ArgsType
+    description: str  # A description of the use case.
+    model: torch.nn.Module
+    name: str
+    example_kwargs: dict[str, Any] = field(default_factory=dict)
+    extra_args: ArgsType | None = None  # For testing graph generalization.
+    # Tags associated with the use case. (e.g dynamic-shape, escape-hatch)
+    tags: set[str] = field(default_factory=set)
+    support_level: SupportLevel = SupportLevel.SUPPORTED
+    dynamic_shapes: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        check_inputs_type(self.example_args, self.example_kwargs)
+        if self.extra_args is not None:
+            check_inputs_type(self.extra_args, {})
+
+        for tag in self.tags:
+            _validate_tag(tag)
+
+        if not isinstance(self.description, str) or len(self.description) == 0:
+            raise ValueError(f'Invalid description: "{self.description}"')
+
+
+_EXAMPLE_CASES: dict[str, ExportCase] = {}
+_MODULES: set[ModuleType] = set()
+_EXAMPLE_CONFLICT_CASES: dict[str, list[ExportCase]] = {}
+_EXAMPLE_REWRITE_CASES: dict[str, list[ExportCase]] = {}
+
+
+def register_db_case(case: ExportCase) -> None:
+    """
+    Registers a user provided ExportCase into example bank.
+    """
+    if case.name in _EXAMPLE_CASES:
+        if case.name not in _EXAMPLE_CONFLICT_CASES:
+            _EXAMPLE_CONFLICT_CASES[case.name] = [_EXAMPLE_CASES[case.name]]
+        _EXAMPLE_CONFLICT_CASES[case.name].append(case)
+        return
+
+    _EXAMPLE_CASES[case.name] = case
+
+
+def to_snake_case(name):
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()
+
+
+def _make_export_case(m, name, configs):
+    if not isinstance(m, torch.nn.Module):
+        raise TypeError("Export case class should be a torch.nn.Module.")
+
+    if "description" not in configs:
+        # Fallback to docstring if description is missing.
+        if m.__doc__ is None:
+            raise AssertionError(
+                f"Could not find description or docstring for export case: {m}"
+            )
+        configs = {**configs, "description": m.__doc__}
+    # pyrefly: ignore [bad-argument-type]
+    return ExportCase(**{**configs, "model": m, "name": name})
+
+
+def export_case(**kwargs):
+    """
+    Decorator for registering a user provided case into example bank.
+    """
+
+    def wrapper(m):
+        configs = kwargs
+        module = inspect.getmodule(m)
+        if module in _MODULES:
+            raise RuntimeError("export_case should only be used once per example file.")
+
+        if module is None:
+            raise AssertionError("module must not be None")
+        _MODULES.add(module)
+        module_name = module.__name__.split(".")[-1]
+        case = _make_export_case(m, module_name, configs)
+        register_db_case(case)
+        return case
+
+    return wrapper
+
+
+def export_rewrite_case(**kwargs):
+    def wrapper(m):
+        configs = kwargs
+
+        parent = configs.pop("parent")
+        if not isinstance(parent, ExportCase):
+            raise AssertionError(f"expected ExportCase, got {type(parent)}")
+        key = parent.name
+        if key not in _EXAMPLE_REWRITE_CASES:
+            _EXAMPLE_REWRITE_CASES[key] = []
+
+        configs["example_args"] = parent.example_args
+        case = _make_export_case(m, to_snake_case(m.__name__), configs)
+        _EXAMPLE_REWRITE_CASES[key].append(case)
+        return case
+
+    return wrapper

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/__init__.py

@@ -0,0 +1,61 @@
+# mypy: allow-untyped-defs
+import dataclasses
+import glob
+import inspect
+from os.path import basename, dirname, isfile, join
+
+import torch
+from torch._export.db.case import (
+    _EXAMPLE_CASES,
+    _EXAMPLE_CONFLICT_CASES,
+    _EXAMPLE_REWRITE_CASES,
+    SupportLevel,
+    export_case,
+    ExportCase,
+)
+
+
+def _collect_examples():
+    case_names = glob.glob(join(dirname(__file__), "*.py"))
+    case_names = [
+        basename(f)[:-3] for f in case_names if isfile(f) and not f.endswith("__init__.py")
+    ]
+
+    case_fields = {f.name for f in dataclasses.fields(ExportCase)}
+    for case_name in case_names:
+        case = __import__(case_name, globals(), locals(), [], 1)
+        variables = [name for name in dir(case) if name in case_fields]
+        export_case(**{v: getattr(case, v) for v in variables})(case.model)
+
+_collect_examples()
+
+def all_examples():
+    return _EXAMPLE_CASES
+
+
+if len(_EXAMPLE_CONFLICT_CASES) > 0:
+
+    def get_name(case):
+        model = case.model
+        if isinstance(model, torch.nn.Module):
+            model = type(model)
+        return model.__name__
+
+    msg = "Error on conflict export case name.\n"
+    for case_name, cases in _EXAMPLE_CONFLICT_CASES.items():
+        msg += f"Case name {case_name} is associated with multiple cases:\n  "
+        msg += f"[{','.join(map(get_name, cases))}]\n"
+
+    raise RuntimeError(msg)
+
+
+def filter_examples_by_support_level(support_level: SupportLevel):
+    return {
+        key: val
+        for key, val in all_examples().items()
+        if val.support_level == support_level
+    }
+
+
+def get_rewrite_cases(case):
+    return _EXAMPLE_REWRITE_CASES.get(case.name, [])

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/assume_constant_result.py

@@ -0,0 +1,20 @@
+# mypy: allow-untyped-defs
+import torch
+import torch._dynamo as torchdynamo
+
+
+class AssumeConstantResult(torch.nn.Module):
+    """
+    Applying `assume_constant_result` decorator to burn make non-tracable code as constant.
+    """
+
+    @torchdynamo.assume_constant_result
+    def get_item(self, y):
+        return y.int().item()
+
+    def forward(self, x, y):
+        return x[: self.get_item(y)]
+
+example_args = (torch.randn(3, 2), torch.tensor(4))
+tags = {"torch.escape-hatch"}
+model = AssumeConstantResult()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/autograd_function.py

@@ -0,0 +1,25 @@
+# mypy: allow-untyped-defs
+import torch
+
+class MyAutogradFunction(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, x):
+        return x.clone()
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        return grad_output + 1
+
+class AutogradFunction(torch.nn.Module):
+    """
+    TorchDynamo does not keep track of backward() on autograd functions. We recommend to
+    use `allow_in_graph` to mitigate this problem.
+    """
+
+    def forward(self, x):
+        return MyAutogradFunction.apply(x)
+
+example_args = (torch.randn(3, 2),)
+model = AutogradFunction()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/class_method.py

@@ -0,0 +1,22 @@
+# mypy: allow-untyped-defs
+import torch
+
+class ClassMethod(torch.nn.Module):
+    """
+    Class methods are inlined during tracing.
+    """
+
+    @classmethod
+    def method(cls, x):
+        return x + 1
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(4, 2)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return self.method(x) * self.__class__.method(x) * type(self).method(x)
+
+example_args = (torch.randn(3, 4),)
+model = ClassMethod()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_class_method.py

@@ -0,0 +1,44 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import cond
+
+class MySubModule(torch.nn.Module):
+    def foo(self, x):
+        return x.cos()
+
+    def forward(self, x):
+        return self.foo(x)
+
+class CondBranchClassMethod(torch.nn.Module):
+    """
+    The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules:
+      - both branches must take the same args, which must also match the branch args passed to cond.
+      - both branches must return a single tensor
+      - returned tensor must have the same tensor metadata, e.g. shape and dtype
+      - branch function can be free function, nested function, lambda, class methods
+      - branch function can not have closure variables
+      - no inplace mutations on inputs or global variables
+
+
+    This example demonstrates using class method in cond().
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.subm = MySubModule()
+
+    def bar(self, x):
+        return x.sin()
+
+    def forward(self, x):
+        return cond(x.shape[0] <= 2, self.subm.forward, self.bar, [x])
+
+example_args = (torch.randn(3),)
+tags = {
+    "torch.cond",
+    "torch.dynamic-shape",
+}
+model = CondBranchClassMethod()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nested_function.py

@@ -0,0 +1,41 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import cond
+
+class CondBranchNestedFunction(torch.nn.Module):
+    """
+    The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules:
+      - both branches must take the same args, which must also match the branch args passed to cond.
+      - both branches must return a single tensor
+      - returned tensor must have the same tensor metadata, e.g. shape and dtype
+      - branch function can be free function, nested function, lambda, class methods
+      - branch function can not have closure variables
+      - no inplace mutations on inputs or global variables
+
+    This example demonstrates using nested function in cond().
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def forward(self, x):
+        def true_fn(x):
+            def inner_true_fn(y):
+                return x + y
+
+            return inner_true_fn(x)
+
+        def false_fn(x):
+            def inner_false_fn(y):
+                return x - y
+
+            return inner_false_fn(x)
+
+        return cond(x.shape[0] < 10, true_fn, false_fn, [x])
+
+example_args = (torch.randn(3),)
+tags = {
+    "torch.cond",
+    "torch.dynamic-shape",
+}
+model = CondBranchNestedFunction()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nonlocal_variables.py

@@ -0,0 +1,59 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import cond
+
+class CondBranchNonlocalVariables(torch.nn.Module):
+    """
+    The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules:
+    - both branches must take the same args, which must also match the branch args passed to cond.
+    - both branches must return a single tensor
+    - returned tensor must have the same tensor metadata, e.g. shape and dtype
+    - branch function can be free function, nested function, lambda, class methods
+    - branch function can not have closure variables
+    - no inplace mutations on inputs or global variables
+
+    This example demonstrates how to rewrite code to avoid capturing closure variables in branch functions.
+
+    The code below will not work because capturing closure variables is not supported.
+    ```
+    my_tensor_var = x + 100
+    my_primitive_var = 3.14
+
+    def true_fn(y):
+        nonlocal my_tensor_var, my_primitive_var
+        return y + my_tensor_var + my_primitive_var
+
+    def false_fn(y):
+        nonlocal my_tensor_var, my_primitive_var
+        return y - my_tensor_var - my_primitive_var
+
+    return cond(x.shape[0] > 5, true_fn, false_fn, [x])
+    ```
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def forward(self, x):
+        my_tensor_var = x + 100
+        my_primitive_var = 3.14
+
+        def true_fn(x, y, z):
+            return x + y + z
+
+        def false_fn(x, y, z):
+            return x - y - z
+
+        return cond(
+            x.shape[0] > 5,
+            true_fn,
+            false_fn,
+            [x, my_tensor_var, torch.tensor(my_primitive_var)],
+        )
+
+example_args = (torch.randn(6),)
+tags = {
+    "torch.cond",
+    "torch.dynamic-shape",
+}
+model = CondBranchNonlocalVariables()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_closed_over_variable.py

@@ -0,0 +1,22 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import cond
+
+class CondClosedOverVariable(torch.nn.Module):
+    """
+    torch.cond() supports branches closed over arbitrary variables.
+    """
+
+    def forward(self, pred, x):
+        def true_fn(val):
+            return x * 2
+
+        def false_fn(val):
+            return x - 2
+
+        return cond(pred, true_fn, false_fn, [x + 1])
+
+example_args = (torch.tensor(True), torch.randn(3, 2))
+tags = {"torch.cond", "python.closure"}
+model = CondClosedOverVariable()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_operands.py

@@ -0,0 +1,35 @@
+# mypy: allow-untyped-defs
+import torch
+
+from torch.export import Dim
+
+x = torch.randn(3, 2)
+y = torch.randn(2)
+dim0_x = Dim("dim0_x")
+
+class CondOperands(torch.nn.Module):
+    """
+    The operands passed to cond() must be:
+    - a list of tensors
+    - match arguments of `true_fn` and `false_fn`
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def forward(self, x, y):
+        def true_fn(x, y):
+            return x + y
+
+        def false_fn(x, y):
+            return x - y
+
+        return torch.cond(x.shape[0] > 2, true_fn, false_fn, [x, y])
+
+example_args = (x, y)
+tags = {
+    "torch.cond",
+    "torch.dynamic-shape",
+}
+extra_inputs = (torch.randn(2, 2), torch.randn(2))
+dynamic_shapes = {"x": {0: dim0_x}, "y": None}
+model = CondOperands()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/cond_predicate.py

@@ -0,0 +1,25 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import cond
+
+class CondPredicate(torch.nn.Module):
+    """
+    The conditional statement (aka predicate) passed to cond() must be one of the following:
+      - torch.Tensor with a single element
+      - boolean expression
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def forward(self, x):
+        pred = x.dim() > 2 and x.shape[2] > 10
+
+        return cond(pred, lambda x: x.cos(), lambda y: y.sin(), [x])
+
+example_args = (torch.randn(6, 4, 3),)
+tags = {
+    "torch.cond",
+    "torch.dynamic-shape",
+}
+model = CondPredicate()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_size_example.py

@@ -0,0 +1,23 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class ConstrainAsSizeExample(torch.nn.Module):
+    """
+    If the value is not known at tracing time, you can provide hint so that we
+    can trace further. Please look at torch._check APIs.
+    """
+
+    def forward(self, x):
+        a = x.item()
+        torch._check(a >= 0)
+        torch._check(a <= 5)
+        return torch.zeros((a, 5))
+
+
+example_args = (torch.tensor(4),)
+tags = {
+    "torch.dynamic-value",
+    "torch.escape-hatch",
+}
+model = ConstrainAsSizeExample()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_value_example.py

@@ -0,0 +1,26 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class ConstrainAsValueExample(torch.nn.Module):
+    """
+    If the value is not known at tracing time, you can provide hint so that we
+    can trace further. Please look at torch._check API.
+    """
+
+    def forward(self, x, y):
+        a = x.item()
+        torch._check(a >= 0)
+        torch._check(a <= 5)
+
+        if a < 6:
+            return y.sin()
+        return y.cos()
+
+
+example_args = (torch.tensor(4), torch.randn(5, 5))
+tags = {
+    "torch.dynamic-value",
+    "torch.escape-hatch",
+}
+model = ConstrainAsValueExample()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/decorator.py

@@ -0,0 +1,23 @@
+# mypy: allow-untyped-defs
+import functools
+
+import torch
+
+def test_decorator(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs) + 1
+
+    return wrapper
+
+class Decorator(torch.nn.Module):
+    """
+    Decorators calls are inlined into the exported function during tracing.
+    """
+
+    @test_decorator
+    def forward(self, x, y):
+        return x + y
+
+example_args = (torch.randn(3, 2), torch.randn(3, 2))
+model = Decorator()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dictionary.py

@@ -0,0 +1,17 @@
+# mypy: allow-untyped-defs
+import torch
+
+class Dictionary(torch.nn.Module):
+    """
+    Dictionary structures are inlined and flattened along tracing.
+    """
+
+    def forward(self, x, y):
+        elements = {}
+        elements["x2"] = x * x
+        y = y * elements["x2"]
+        return {"y": y}
+
+example_args = (torch.randn(3, 2), torch.tensor(4))
+tags = {"python.data-structure"}
+model = Dictionary()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_assert.py

@@ -0,0 +1,18 @@
+# mypy: allow-untyped-defs
+import torch
+
+class DynamicShapeAssert(torch.nn.Module):
+    """
+    A basic usage of python assertion.
+    """
+
+    def forward(self, x):
+        # assertion with error message
+        assert x.shape[0] > 2, f"{x.shape[0]} is greater than 2"  # noqa: S101
+        # assertion without error message
+        assert x.shape[0] > 1  # noqa: S101
+        return x
+
+example_args = (torch.randn(3, 2),)
+tags = {"python.assert"}
+model = DynamicShapeAssert()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_constructor.py

@@ -0,0 +1,15 @@
+# mypy: allow-untyped-defs
+import torch
+
+class DynamicShapeConstructor(torch.nn.Module):
+    """
+    Tensor constructors should be captured with dynamic shape inputs rather
+    than being baked in with static shape.
+    """
+
+    def forward(self, x):
+        return torch.zeros(x.shape[0] * 2)
+
+example_args = (torch.randn(3, 2),)
+tags = {"torch.dynamic-shape"}
+model = DynamicShapeConstructor()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_if_guard.py

@@ -0,0 +1,19 @@
+# mypy: allow-untyped-defs
+import torch
+
+class DynamicShapeIfGuard(torch.nn.Module):
+    """
+    `if` statement with backed dynamic shape predicate will be specialized into
+    one particular branch and generate a guard. However, export will fail if the
+    the dimension is marked as dynamic shape from higher level API.
+    """
+
+    def forward(self, x):
+        if x.shape[0] == 3:
+            return x.cos()
+
+        return x.sin()
+
+example_args = (torch.randn(3, 2, 2),)
+tags = {"torch.dynamic-shape", "python.control-flow"}
+model = DynamicShapeIfGuard()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_map.py

@@ -0,0 +1,19 @@
+# mypy: allow-untyped-defs
+import torch
+
+from functorch.experimental.control_flow import map
+
+class DynamicShapeMap(torch.nn.Module):
+    """
+    functorch map() maps a function over the first tensor dimension.
+    """
+
+    def forward(self, xs, y):
+        def body(x, y):
+            return x + y
+
+        return map(body, xs, y)
+
+example_args = (torch.randn(3, 2), torch.randn(2))
+tags = {"torch.dynamic-shape", "torch.map"}
+model = DynamicShapeMap()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_round.py

@@ -0,0 +1,21 @@
+# mypy: allow-untyped-defs
+import torch
+
+from torch._export.db.case import SupportLevel
+from torch.export import Dim
+
+class DynamicShapeRound(torch.nn.Module):
+    """
+    Calling round on dynamic shapes is not supported.
+    """
+
+    def forward(self, x):
+        return x[: round(x.shape[0] / 2)]
+
+x = torch.randn(3, 2)
+dim0_x = Dim("dim0_x")
+example_args = (x,)
+tags = {"torch.dynamic-shape", "python.builtin"}
+support_level = SupportLevel.NOT_SUPPORTED_YET
+dynamic_shapes = {"x": {0: dim0_x}}
+model = DynamicShapeRound()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_slicing.py

@@ -0,0 +1,15 @@
+# mypy: allow-untyped-defs
+import torch
+
+class DynamicShapeSlicing(torch.nn.Module):
+    """
+    Slices with dynamic shape arguments should be captured into the graph
+    rather than being baked in.
+    """
+
+    def forward(self, x):
+        return x[: x.shape[0] - 2, x.shape[1] - 1 :: 2]
+
+example_args = (torch.randn(3, 2),)
+tags = {"torch.dynamic-shape"}
+model = DynamicShapeSlicing()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_view.py

@@ -0,0 +1,17 @@
+# mypy: allow-untyped-defs
+import torch
+
+class DynamicShapeView(torch.nn.Module):
+    """
+    Dynamic shapes should be propagated to view arguments instead of being
+    baked into the exported graph.
+    """
+
+    def forward(self, x):
+        new_x_shape = x.size()[:-1] + (2, 5)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1)
+
+example_args = (torch.randn(10, 10),)
+tags = {"torch.dynamic-shape"}
+model = DynamicShapeView()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/fn_with_kwargs.py

@@ -0,0 +1,30 @@
+# mypy: allow-untyped-defs
+import torch
+
+class FnWithKwargs(torch.nn.Module):
+    """
+    Keyword arguments are not supported at the moment.
+    """
+
+    def forward(self, pos0, tuple0, *myargs, mykw0, **mykwargs):
+        out = pos0
+        for arg in tuple0:
+            out = out * arg
+        for arg in myargs:
+            out = out * arg
+        out = out * mykw0
+        out = out * mykwargs["input0"] * mykwargs["input1"]
+        return out
+
+example_args = (
+    torch.randn(4),
+    (torch.randn(4), torch.randn(4)),
+    *[torch.randn(4), torch.randn(4)]
+)
+example_kwargs = {
+    "mykw0": torch.randn(4),
+    "input0": torch.randn(4),
+    "input1": torch.randn(4),
+}
+tags = {"python.data-structure"}
+model = FnWithKwargs()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/list_contains.py

@@ -0,0 +1,17 @@
+# mypy: allow-untyped-defs
+import torch
+
+class ListContains(torch.nn.Module):
+    """
+    List containment relation can be checked on a dynamic shape or constants.
+    """
+
+    def forward(self, x):
+        assert x.size(-1) in [6, 2]  # noqa: S101
+        assert x.size(0) not in [4, 5, 6]  # noqa: S101
+        assert "monkey" not in ["cow", "pig"]  # noqa: S101
+        return x + x
+
+example_args = (torch.randn(3, 2),)
+tags = {"torch.dynamic-shape", "python.data-structure", "python.assert"}
+model = ListContains()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/list_unpack.py

@@ -0,0 +1,21 @@
+# mypy: allow-untyped-defs
+
+import torch
+
+class ListUnpack(torch.nn.Module):
+    """
+    Lists are treated as static construct, therefore unpacking should be
+    erased after tracing.
+    """
+
+    def forward(self, args: list[torch.Tensor]):
+        """
+        Lists are treated as static construct, therefore unpacking should be
+        erased after tracing.
+        """
+        x, *y = args
+        return x + y[0]
+
+example_args = ([torch.randn(3, 2), torch.tensor(4), torch.tensor(5)],)
+tags = {"python.control-flow", "python.data-structure"}
+model = ListUnpack()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/model_attr_mutation.py

@@ -0,0 +1,24 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class ModelAttrMutation(torch.nn.Module):
+    """
+    Attribute mutation raises a warning. Covered in the test_export.py test_detect_leak_strict test.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.attr_list = [torch.randn(3, 2), torch.randn(3, 2)]
+
+    def recreate_list(self):
+        return [torch.zeros(3, 2), torch.zeros(3, 2)]
+
+    def forward(self, x):
+        self.attr_list = self.recreate_list()
+        return x.sum() + self.attr_list[0].sum()
+
+
+example_args = (torch.randn(3, 2),)
+tags = {"python.object-model"}
+model = ModelAttrMutation()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/nested_function.py

@@ -0,0 +1,23 @@
+# mypy: allow-untyped-defs
+import torch
+
+class NestedFunction(torch.nn.Module):
+    """
+    Nested functions are traced through. Side effects on global captures
+    are not supported though.
+    """
+
+    def forward(self, a, b):
+        x = a + b
+        z = a - b
+
+        def closure(y):
+            nonlocal x
+            x += 1
+            return x * y + z
+
+        return closure(x)
+
+example_args = (torch.randn(3, 2), torch.randn(2))
+tags = {"python.closure"}
+model = NestedFunction()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/null_context_manager.py

@@ -0,0 +1,21 @@
+# mypy: allow-untyped-defs
+import contextlib
+
+import torch
+
+class NullContextManager(torch.nn.Module):
+    """
+    Null context manager in Python will be traced out.
+    """
+
+    def forward(self, x):
+        """
+        Null context manager in Python will be traced out.
+        """
+        ctx = contextlib.nullcontext()
+        with ctx:
+            return x.sin() + x.cos()
+
+example_args = (torch.randn(3, 2),)
+tags = {"python.context-manager"}
+model = NullContextManager()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/optional_input.py

@@ -0,0 +1,20 @@
+# mypy: allow-untyped-defs
+import torch
+from torch._export.db.case import SupportLevel
+
+
+class OptionalInput(torch.nn.Module):
+    """
+    Tracing through optional input is not supported yet
+    """
+
+    def forward(self, x, y=torch.randn(2, 3)):
+        if y is not None:
+            return x + y
+        return x
+
+
+example_args = (torch.randn(2, 3),)
+tags = {"python.object-model"}
+support_level = SupportLevel.SUPPORTED
+model = OptionalInput()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/pytree_flatten.py

@@ -0,0 +1,16 @@
+# mypy: allow-untyped-defs
+import torch
+
+from torch.utils import _pytree as pytree
+
+class PytreeFlatten(torch.nn.Module):
+    """
+    Pytree from PyTorch can be captured by TorchDynamo.
+    """
+
+    def forward(self, x):
+        y, _spec = pytree.tree_flatten(x)
+        return y[0] + 1
+
+example_args = ({1: torch.randn(3, 2), 2: torch.randn(3, 2)},),
+model = PytreeFlatten()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/scalar_output.py

@@ -0,0 +1,23 @@
+# mypy: allow-untyped-defs
+import torch
+
+from torch.export import Dim
+
+x = torch.randn(3, 2)
+dim1_x = Dim("dim1_x")
+
+class ScalarOutput(torch.nn.Module):
+    """
+    Returning scalar values from the graph is supported, in addition to Tensor
+    outputs. Symbolic shapes are captured and rank is specialized.
+    """
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x):
+        return x.shape[1] + 1
+
+example_args = (x,)
+tags = {"torch.dynamic-shape"}
+dynamic_shapes = {"x": {1: dim1_x}}
+model = ScalarOutput()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/specialized_attribute.py

@@ -0,0 +1,26 @@
+# mypy: allow-untyped-defs
+from enum import Enum
+
+import torch
+
+class Animal(Enum):
+    COW = "moo"
+
+class SpecializedAttribute(torch.nn.Module):
+    """
+    Model attributes are specialized.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.a = "moo"
+        self.b = 4
+
+    def forward(self, x):
+        if self.a == Animal.COW.value:
+            return x * x + self.b
+        else:
+            raise ValueError("bad")
+
+example_args = (torch.randn(3, 2),)
+model = SpecializedAttribute()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/static_for_loop.py

@@ -0,0 +1,16 @@
+# mypy: allow-untyped-defs
+import torch
+
+class StaticForLoop(torch.nn.Module):
+    """
+    A for loop with constant number of iterations should be unrolled in the exported graph.
+    """
+
+    def forward(self, x):
+        # constant
+        ret = [i + x for i in range(10)]
+        return ret
+
+example_args = (torch.randn(3, 2),)
+tags = {"python.control-flow"}
+model = StaticForLoop()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/static_if.py

@@ -0,0 +1,18 @@
+# mypy: allow-untyped-defs
+import torch
+
+class StaticIf(torch.nn.Module):
+    """
+    `if` statement with static predicate value should be traced through with the
+    taken branch.
+    """
+
+    def forward(self, x):
+        if len(x.shape) == 3:
+            return x + torch.ones(1, 1, 1)
+
+        return x
+
+example_args = (torch.randn(3, 2, 2),)
+tags = {"python.control-flow"}
+model = StaticIf()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/tensor_setattr.py

@@ -0,0 +1,15 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class TensorSetattr(torch.nn.Module):
+    """
+    setattr() call onto tensors is not supported.
+    """
+    def forward(self, x, attr):
+        setattr(x, attr, torch.randn(3, 2))
+        return x + 4
+
+example_args = (torch.randn(3, 2), "attr")
+tags = {"python.builtin"}
+model = TensorSetattr()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/type_reflection_method.py

@@ -0,0 +1,22 @@
+# mypy: allow-untyped-defs
+import torch
+
+class A:
+    @classmethod
+    def func(cls, x):
+        return 1 + x
+
+class TypeReflectionMethod(torch.nn.Module):
+    """
+    type() calls on custom objects followed by attribute accesses are not allowed
+    due to its overly dynamic nature.
+    """
+
+    def forward(self, x):
+        a = A()
+        return type(a).func(x)
+
+
+example_args = (torch.randn(3, 4),)
+tags = {"python.builtin"}
+model = TypeReflectionMethod()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/unsupported_operator.py

@@ -0,0 +1,18 @@
+# mypy: allow-untyped-defs
+import torch
+from torch._export.db.case import SupportLevel
+
+
+class TorchSymMin(torch.nn.Module):
+    """
+    torch.sym_min operator is supported in export.
+    """
+
+    def forward(self, x):
+        return x.sum() + torch.sym_min(x.size(0), 100)
+
+
+example_args = (torch.randn(3, 2),)
+tags = {"torch.operator"}
+support_level = SupportLevel.SUPPORTED
+model = TorchSymMin()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/examples/user_input_mutation.py

@@ -0,0 +1,17 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class UserInputMutation(torch.nn.Module):
+    """
+    Directly mutate user input in forward
+    """
+
+    def forward(self, x):
+        x.mul_(2)
+        return x.cos()
+
+
+example_args = (torch.randn(3, 2),)
+tags = {"torch.mutation"}
+model = UserInputMutation()

kronos-venv/lib/python3.12/site-packages/torch/_export/db/gen_example.py

@@ -0,0 +1,23 @@
+import os
+import sys
+
+import torch._export.db.examples as examples
+
+TEMPLATE = '''import torch
+
+def {case_name}(x):
+    """
+    """
+
+    return
+'''
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        raise AssertionError(f"expected 2 arguments, got {len(sys.argv)}")
+    root_dir = examples.__name__.replace(".", "/")
+    if not os.path.exists(root_dir):
+        raise AssertionError(f"root_dir does not exist: {root_dir}")
+    with open(os.path.join(root_dir, sys.argv[1] + ".py"), "w") as f:
+        print("Writing to", f.name, "...")
+        f.write(TEMPLATE.format(case_name=sys.argv[1]))

kronos-venv/lib/python3.12/site-packages/torch/_export/db/logging.py

@@ -0,0 +1,46 @@
+
+def exportdb_error_message(case_name: str) -> str:
+    from .examples import all_examples
+    from torch._utils_internal import log_export_usage
+
+    ALL_EXAMPLES = all_examples()
+    # Detect whether case_name is really registered in exportdb.
+    if case_name in ALL_EXAMPLES:
+        url_case_name = case_name.replace("_", "-")
+        return f"See {case_name} in exportdb for unsupported case. \
+                https://pytorch.org/docs/main/generated/exportdb/index.html#{url_case_name}"
+    else:
+        log_export_usage(
+            event="export.error.casenotregistered",
+            message=case_name,
+        )
+        return f"{case_name} is unsupported."
+
+
+def get_class_if_classified_error(e: Exception) -> str | None:
+    """
+    Returns a string case name if the export error e is classified.
+    Returns None otherwise.
+    """
+
+    from torch._dynamo.exc import TorchRuntimeError, Unsupported, UserError
+
+    ALWAYS_CLASSIFIED = "always_classified"
+    DEFAULT_CLASS_SIGIL = "case_name"
+
+    # add error types that should be classified, along with any attribute name
+    # whose presence acts like a sigil to further distinguish which errors of
+    # that type should be classified. If the attribute name is None, then the
+    # error type is always classified.
+    _ALLOW_LIST = {
+        Unsupported: DEFAULT_CLASS_SIGIL,
+        UserError: DEFAULT_CLASS_SIGIL,
+        TorchRuntimeError: None,
+    }
+    if type(e) in _ALLOW_LIST:
+        # pyrefly: ignore [bad-index, index-error]
+        attr_name = _ALLOW_LIST[type(e)]
+        if attr_name is None:
+            return ALWAYS_CLASSIFIED
+        return getattr(e, attr_name, None)
+    return None

kronos-venv/lib/python3.12/site-packages/torch/_export/error.py

@@ -0,0 +1,56 @@
+from enum import Enum
+
+
+class ExportErrorType(Enum):
+    # User providing invalid inputs to either tracer, or other public facing APIs
+    INVALID_INPUT_TYPE = 1
+
+    # User returning values from their models that we don't support.
+    INVALID_OUTPUT_TYPE = 2
+
+    # Generated IR does not conform to Export IR Specification.
+    VIOLATION_OF_SPEC = 3
+
+    # User's code contains types and functionalities we don't support.
+    NOT_SUPPORTED = 4
+
+    # User's code didn't provide necessary details for us to successfully trace and export.
+    # For example, we use a lot of decorators and ask users to annotate their model.
+    MISSING_PROPERTY = 5
+
+    # User is using an API without proper initialization step.
+    UNINITIALIZED = 6
+
+
+def internal_assert(pred: bool, assert_msg: str) -> None:
+    """
+    This is exir's custom assert method. It internally just throws InternalError.
+    Note that the sole purpose is to throw our own error while maintaining similar syntax
+    as python assert.
+    """
+
+    if not pred:
+        raise InternalError(assert_msg)
+
+
+class InternalError(Exception):
+    """
+    Raised when an internal invariance is violated in EXIR stack.
+    Should hint users to report a bug to dev and expose the original
+    error message.
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class ExportError(Exception):
+    """
+    This type of exception is raised for errors that are directly caused by the user
+    code. In general, user errors happen during model authoring, tracing, using our public
+    facing APIs, and writing graph passes.
+    """
+
+    def __init__(self, error_code: ExportErrorType, message: str) -> None:
+        prefix = f"[{error_code}]: "
+        super().__init__(prefix + message)

kronos-venv/lib/python3.12/site-packages/torch/_export/pass_base.py

@@ -0,0 +1,497 @@
+# mypy: allow-untyped-defs
+import operator
+import traceback
+from collections.abc import Callable
+from contextlib import nullcontext
+from typing import Any
+
+import torch
+from torch import fx
+from torch._dispatch.python import enable_python_dispatcher
+from torch._export.pass_infra.node_metadata import NodeMetadata
+from torch._export.pass_infra.proxy_value import ProxyValue
+from torch._higher_order_ops.map import _unstack_pytree
+from torch._subclasses import FakeTensor, UnsupportedFakeTensorException
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx import traceback as fx_traceback
+from torch.fx.experimental.proxy_tensor import PythonKeyTracer
+from torch.fx.experimental.symbolic_shapes import (
+    compute_unbacked_bindings,
+    PropagateUnbackedSymInts,
+)
+from torch.fx.graph import CodeGen
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from torch.fx.passes.shape_prop import _extract_tensor_metadata, TensorMetadata
+from torch.utils import _pytree as pytree
+
+
+__all__ = ["_ExportPassBaseDeprecatedDoNotUse"]
+
+
+Argument = Any
+Value = Any
+Fn = Callable[..., Any]
+PassType = Callable[[torch.fx.GraphModule], PassResult | None]
+
+
+_TORCH_SYM_OPS: set[Callable] = {
+    torch.sym_int,
+    torch.sym_float,
+    torch.sym_ite,
+    torch.sym_max,
+    torch.sym_min,
+    torch.sym_not,
+    torch.sym_sqrt,
+}
+
+
+class ExportPassBaseError(RuntimeError):
+    pass
+
+
+class _ExportPassBaseDeprecatedDoNotUse(PassBase):
+    """
+    Interpreter-based pass class to help users maintain the IR spec while writing
+    transformations.
+    """
+
+    @staticmethod
+    def _create_dummy_node_metadata():
+        return NodeMetadata({"stack_trace": "".join(traceback.format_stack(limit=1))})
+
+    class ExportTracer(PythonKeyTracer):
+        def __init__(
+            self, callback: "_ExportPassBaseDeprecatedDoNotUse", codegen: CodeGen
+        ) -> None:
+            super().__init__()
+            self.callback = callback
+            self.root = torch.nn.Module()
+            self.graph = torch.fx.Graph()
+            self.graph.set_codegen(codegen)
+            self.tensor_attrs: dict[str, torch.Tensor] = {}  # type: ignore[assignment]
+            self.fake_tensor_mode: FakeTensorMode | None = None
+            self.submodules: dict[torch.nn.Module, str] = {}
+
+        def trace(self) -> None:  # type: ignore[override]
+            raise ExportPassBaseError("ExportTracer doesn't support trace().")
+
+        def create_arg(self, a: Argument) -> torch.fx.Node:
+            if isinstance(a, torch.nn.Module):
+                if a not in self.submodules:
+                    name_submodule = f"submodule_{len(self.submodules)}"
+                    self.root.add_module(name_submodule, a)
+                    self.submodules[a] = name_submodule
+            elif isinstance(a, FakeTensor):
+                if not hasattr(a, "constant") or a.constant is None:
+                    raise ExportPassBaseError(f"Cannot add {a} to graph.")
+                a = a.constant
+            node = super().create_arg(a)
+            if (
+                isinstance(a, torch.Tensor)
+                and isinstance(node, torch.fx.Node)
+                and node.op == "get_attr"
+            ):
+                self.set_metadata(node, a)
+                self.callback.on_attr(ProxyValue(a, node))
+            return node
+
+        def set_metadata(
+            self,
+            node: torch.fx.Node,
+            value: Argument,
+        ) -> None:
+            # propagate the fake tensor or sym nodes
+            def make_val(
+                x: Argument,
+            ) -> (
+                FakeTensor
+                | torch.SymInt
+                | torch.SymFloat
+                | torch.SymBool
+                | int
+                | float
+                | bool
+                | str
+                | None
+            ):
+                if isinstance(x, FakeTensor):
+                    return x
+                elif isinstance(x, torch.Tensor):
+                    if x.is_quantized:
+                        # TODO (tmanlaibaatar) properly support Quantized FakeTensor
+                        x = torch.dequantize(x)
+
+                    try:
+                        if self.fake_tensor_mode is None:
+                            raise AssertionError("fake_tensor_mode must not be None")
+                        # TODO we should allocate static shapes
+                        # for param/buffer values
+                        if isinstance(x, torch.nn.Parameter):
+                            fake_tensor = self.fake_tensor_mode.from_tensor(
+                                x, static_shapes=True
+                            )
+                        else:
+                            fake_tensor = self.fake_tensor_mode.from_tensor(x)
+                    except UnsupportedFakeTensorException:
+                        # TODO: This is just a workaround to get over the
+                        # x.as_subclass error
+                        print(
+                            "Fakeifying a Tensor subclass is not supported \
+                            right now. Instead a TensorMetadata is used."
+                        )
+                        fake_tensor = None
+                    return fake_tensor
+                elif isinstance(
+                    x,
+                    (
+                        torch.SymInt,
+                        torch.SymFloat,
+                        torch.SymBool,
+                        int,
+                        float,
+                        bool,
+                        str,
+                    ),
+                ):
+                    return x
+                else:
+                    return None
+
+            node.meta["val"] = pytree.tree_map(make_val, value)
+
+            # Set the tensor_metadata for values that do not have a corresponding FakeTensor
+            def make_tensor_meta(x: Argument) -> TensorMetadata | None:
+                if not isinstance(x, FakeTensor) and isinstance(x, torch.Tensor):
+                    if x.is_quantized:
+                        # TODO (tmanlaibaatar) properly support Quantized FakeTensor
+                        x = torch.dequantize(x)
+
+                    try:
+                        if self.fake_tensor_mode is None:
+                            raise AssertionError("fake_tensor_mode must not be None")
+                        _ = self.fake_tensor_mode.from_tensor(x)
+                        tensor_meta = None
+                    except UnsupportedFakeTensorException:
+                        # TODO: This is just a workaround to get over the
+                        # x.as_subclass error
+                        tensor_meta = _extract_tensor_metadata(x)
+                    return tensor_meta
+                else:
+                    return None
+
+            node.meta["tensor_meta"] = pytree.tree_map(make_tensor_meta, value)
+
+    class ExportInterpreter(fx.Interpreter):
+        def __init__(
+            self, callback: "_ExportPassBaseDeprecatedDoNotUse", gm: fx.GraphModule
+        ) -> None:
+            super().__init__(gm)
+            self.callback = callback
+            self.node: torch.fx.Node = next(iter(gm.graph.nodes))
+
+        # pyrefly: ignore [bad-override]
+        def placeholder(
+            self,
+            target: str,  # type: ignore[override]
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> ProxyValue:
+            arg = super().placeholder(target, args, kwargs)
+            return self.callback.placeholder(target, arg, NodeMetadata(self.node.meta))
+
+        def output(
+            self,
+            target: torch.fx.node.Target,
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> ProxyValue:
+            return self.callback.output(args[0], NodeMetadata(self.node.meta)).data  # type: ignore[return-value]
+
+        def call_function(
+            self,
+            target: torch.fx.node.Target,
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> ProxyValue:
+            meta = NodeMetadata(self.node.meta)
+
+            if target is operator.getitem:
+                value, key = args
+                return self.callback.call_getitem(value, key, meta)
+            elif getattr(target, "__module__", None) in {
+                "_operator",
+                "builtins",
+                "math",
+            }:
+                if not callable(target):
+                    raise AssertionError(f"expected callable target, got {target}")
+                return self.callback.call_sym(target, args, meta)
+            elif target in _TORCH_SYM_OPS:
+                if not callable(target):
+                    raise AssertionError(f"expected callable target, got {target}")
+                return self.callback.call_sym(target, args, meta)
+            elif isinstance(
+                target, (torch._ops.OpOverload, torch._ops.OpOverloadPacket)
+            ):
+                return self.callback.call_operator(
+                    target,
+                    args,
+                    kwargs,
+                    meta,
+                )
+            elif target is torch.ops.higher_order.cond:
+                pred, true_fn, false_fn, inputs = args
+                return self.callback.call_cond(pred, true_fn, false_fn, inputs, meta)
+            elif target is torch.ops.higher_order.map_impl:
+                f, mapped_args, operands = args  # type: ignore[assignment]
+                return self.callback.call_map(f, mapped_args, operands, meta)
+            # For other unregistered HigherOrderOps, just interpret them blindly
+            elif isinstance(target, torch._ops.HigherOrderOperator):
+                return self.callback._fx(
+                    "call_function",
+                    target,
+                    args,
+                    kwargs,
+                    meta,
+                )
+            else:
+                raise ExportPassBaseError(f"Unsupported target type: {target}")
+
+        def get_attr(  # type: ignore[override]
+            self,
+            target: str,
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> Argument:
+            return super().get_attr(target, args, kwargs)
+
+        def call_module(
+            self,
+            target: torch.fx.node.Target,
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> None:
+            raise ExportPassBaseError("call_module is not supported.")
+
+        def call_method(  # type: ignore[override]
+            self,
+            target: str,
+            args: tuple[Argument, ...],
+            kwargs: dict[str, Argument],
+        ) -> None:
+            raise ExportPassBaseError("call_method is not supported.")
+
+        def run_node(self, n: torch.fx.Node) -> Argument:
+            self.node = n
+            self.callback.node_debug_str = n.format_node()
+            return super().run_node(n)
+
+    def __init__(self) -> None:
+        self.interpreter = PropagateUnbackedSymInts(
+            torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
+        )
+        self.tracer = self.ExportTracer(self, CodeGen())
+        self.fake_tensor_mode: FakeTensorMode | None = None
+        self._initialized = True
+        self.node_debug_str: str | None = None
+
+    def _fx(
+        self,
+        kind: str,
+        target: torch.fx.node.Target,
+        args: tuple[Argument, ...],
+        kwargs: dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        args_data, kwargs_data = pytree.tree_map_only(
+            ProxyValue, lambda x: x.data, (args, kwargs)
+        )
+        res_data = getattr(self.interpreter, kind)(target, args_data, kwargs_data)
+        args_proxy, kwargs_proxy = pytree.tree_map_only(
+            ProxyValue, lambda x: x.proxy, (args, kwargs)
+        )
+
+        name = None
+        if isinstance(target, torch._ops.OpOverload):
+            name = self.tracer.graph._target_to_str(target.overloadpacket.__name__)
+
+        res_proxy = self.tracer.create_proxy(
+            kind, target, args_proxy, kwargs_proxy, name=name
+        )
+        res_proxy.node.meta.update(meta.data)
+        if self.fake_tensor_mode and (shape_env := self.fake_tensor_mode.shape_env):
+            if symbol_to_path := compute_unbacked_bindings(shape_env, res_data):
+                res_proxy.node.meta["unbacked_bindings"] = symbol_to_path
+        self.tracer.set_metadata(res_proxy.node, res_data)
+        return ProxyValue(res_data, res_proxy)
+
+    def inputs(self, graph_module: torch.fx.GraphModule) -> list[Argument]:
+        # TODO(angelayi): Update this with what we decide to do for metadata in
+        # the exported graph module
+        if (args := graph_module.meta.get("args", None)) is not None:
+            return list(args)
+
+        def extract_input(node: torch.fx.Node) -> FakeTensor | None:
+            if "val" in node.meta:
+                fake = node.meta["val"]
+                if hasattr(fake, "constant") and fake.constant is not None:
+                    return fake.constant
+                return fake
+            elif tensor_meta := node.meta.get("tensor_meta"):
+                if self.fake_tensor_mode is None:
+                    raise AssertionError("fake_tensor_mode must not be None")
+                return FakeTensor(
+                    self.fake_tensor_mode,
+                    torch.empty(
+                        tensor_meta.shape,
+                        dtype=tensor_meta.dtype,
+                        device="meta",
+                        requires_grad=tensor_meta.requires_grad,
+                        memory_format=tensor_meta.memory_format,
+                    ),
+                    torch.device("cpu"),
+                )
+            elif len(node.users) == 0:
+                return None
+            raise ExportPassBaseError(
+                f"Cannot construct an input for graph module: {graph_module}.",
+            )
+
+        return [
+            extract_input(node)
+            for node in graph_module.graph.nodes
+            if node.op == "placeholder"
+        ]
+
+    def on_attr(self, attr: ProxyValue) -> None:
+        pass
+
+    def placeholder(self, name: str, arg: Argument, meta: NodeMetadata) -> ProxyValue:
+        arg_proxy = self.tracer.create_proxy("placeholder", name, (), {})
+        arg_proxy.node.meta = meta.data
+        self.tracer.set_metadata(arg_proxy.node, arg)
+        return ProxyValue(arg, arg_proxy)
+
+    def call_operator(
+        self,
+        op,
+        args: tuple[Argument, ...],
+        kwargs: dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        return self._fx("call_function", op, args, kwargs, meta)
+
+    def call_sym(
+        self,
+        target: Fn,
+        args: tuple[Argument, ...],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        return self._fx("call_function", target, args, {}, meta)
+
+    def call_cond(
+        self,
+        pred: ProxyValue,
+        true_fn: torch.fx.GraphModule,
+        false_fn: torch.fx.GraphModule,
+        inputs: list[Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        true_branch = self.call_submodule(true_fn, tuple(inputs))
+        false_branch = self.call_submodule(false_fn, tuple(inputs))
+        if true_branch is None:
+            raise AssertionError("true_branch must not be None")
+        if false_branch is None:
+            raise AssertionError("false_branch must not be None")
+        return self._fx(
+            "call_function",
+            torch.ops.higher_order.cond,
+            (pred, true_branch.graph_module, false_branch.graph_module, list(inputs)),
+            {},
+            meta,
+        )
+
+    def call_map(
+        self,
+        f: torch.fx.GraphModule,
+        mapped_args: list[ProxyValue],
+        operands: list[ProxyValue],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        xs = _unstack_pytree([arg.data for arg in mapped_args])[0]
+        f_branch = self.call_submodule(f, tuple(xs + [arg.data for arg in operands]))
+        if f_branch is None:
+            raise AssertionError("f_branch must not be None")
+        return self._fx(
+            "call_function",
+            torch.ops.higher_order.map_impl,
+            (f_branch.graph_module, mapped_args, operands),
+            {},
+            meta,
+        )
+
+    def call_getitem(
+        self, value: ProxyValue, key: int, meta: NodeMetadata
+    ) -> ProxyValue:
+        return self._fx("call_function", operator.getitem, (value, key), {}, meta)
+
+    def output(self, results: list[Argument], meta: NodeMetadata) -> ProxyValue:
+        return self._fx("output", "output", (results,), {}, meta)
+
+    def call_submodule(
+        self, graph_module: fx.GraphModule, inputs: tuple[Argument, ...]
+    ) -> PassResult:
+        prev_tracer, self.tracer = (
+            self.tracer,
+            self.ExportTracer(self, graph_module.graph._codegen),
+        )
+        self.tracer.fake_tensor_mode = prev_tracer.fake_tensor_mode
+        interpreter = self.ExportInterpreter(self, graph_module)
+        # pyrefly: ignore [bad-assignment]
+        prev_interpreter, self.interpreter = (
+            self.interpreter,
+            torch.fx.Interpreter(  # type: ignore[assignment]
+                torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
+            ),
+        )
+        inputs_data = pytree.tree_map_only(ProxyValue, lambda x: x.data, inputs)
+        with fx_traceback.preserve_node_meta():
+            interpreter.run(*inputs_data)
+
+        new_graph_module = torch.fx.GraphModule(self.tracer.root, self.tracer.graph)
+
+        self.tracer = prev_tracer
+        self.interpreter = prev_interpreter
+        return PassResult(
+            new_graph_module,
+            True,
+        )
+
+    def call(self, graph_module: fx.GraphModule) -> PassResult:
+        if not getattr(self, "_initialized", False):
+            raise ExportPassBaseError(
+                "ExportPass is not initialized with __init__().",
+            )
+
+        inputs = self.inputs(graph_module)
+
+        fake_tensor_mode = None
+        for i in inputs:
+            if isinstance(i, FakeTensor):
+                if fake_tensor_mode is not None and fake_tensor_mode is not i.fake_mode:
+                    raise AssertionError("Multiple fake tensor mode detected.")
+                fake_tensor_mode = i.fake_mode
+        if fake_tensor_mode is None:
+            self.tracer.fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True)
+            fake_tensor_mode = nullcontext()  # type: ignore[assignment]
+            dispatcher_mode = nullcontext()  # type: ignore[assignment]
+        else:
+            fake_tensor_mode.allow_non_fake_inputs = True
+            self.tracer.fake_tensor_mode = fake_tensor_mode
+            dispatcher_mode = enable_python_dispatcher()  # type: ignore[assignment]
+        self.fake_tensor_mode = self.tracer.fake_tensor_mode
+
+        with fake_tensor_mode, dispatcher_mode:  # type: ignore[assignment, union-attr]
+            result = self.call_submodule(graph_module, tuple(inputs))
+
+        return result

kronos-venv/lib/python3.12/site-packages/torch/_export/pass_infra/node_metadata.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+
+NodeMetadataValue = Any
+
+
+PROTECTED_KEYS: set[str] = {
+    "val",
+    "stack_trace",
+    "nn_module_stack",
+    "debug_handle",
+    "tensor_meta",
+}
+
+
+class NodeMetadata:
+    def __init__(self, data: dict[str, Any]) -> None:
+        self.data: dict[str, Any] = data.copy()
+
+    def __getitem__(self, key: str) -> NodeMetadataValue:
+        return self.data[key]
+
+    def __setitem__(self, key: str, value: NodeMetadataValue) -> NodeMetadataValue:
+        if key in PROTECTED_KEYS:
+            raise RuntimeError(f"Could not override node key: {key}")
+        self.data[key] = value
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.data
+
+    def copy(self) -> "NodeMetadata":
+        return NodeMetadata(self.data.copy())

kronos-venv/lib/python3.12/site-packages/torch/_export/pass_infra/proxy_value.py

@@ -0,0 +1,49 @@
+# pyre-strict
+from collections.abc import Iterable, Iterator
+from typing import Generic, TypeVar
+
+import torch
+
+
+_T = TypeVar("_T")
+
+
+class ProxyValue(Generic[_T]):
+    # pyre-ignore
+    def __init__(self, data: Iterable[_T], proxy: torch.fx.Proxy | torch.fx.Node):
+        # pyre-ignore
+        self.data = data
+        self.proxy_or_node = proxy
+
+    @property
+    def node(self) -> torch.fx.Node:
+        if isinstance(self.proxy_or_node, torch.fx.Node):
+            return self.proxy_or_node
+        if not isinstance(self.proxy_or_node, torch.fx.Proxy):
+            raise AssertionError(
+                f"expected Node or Proxy, got {type(self.proxy_or_node)}"
+            )
+        return self.proxy_or_node.node
+
+    @property
+    def proxy(self) -> torch.fx.Proxy:
+        if not isinstance(self.proxy_or_node, torch.fx.Proxy):
+            raise RuntimeError(
+                f"ProxyValue doesn't have attached Proxy object. Node: {self.proxy_or_node.format_node()}"
+            )
+        return self.proxy_or_node
+
+    def to_tensor(self) -> torch.Tensor:
+        if not isinstance(self.data, torch.Tensor):
+            raise AssertionError(f"expected Tensor, got {type(self.data)}")
+        return self.data
+
+    def is_tensor(self) -> bool:
+        return isinstance(self.data, torch.Tensor)
+
+    # pyre-ignore
+    def __iter__(self) -> Iterator[_T]:
+        yield from self.data
+
+    def __bool__(self) -> bool:
+        return bool(self.data)

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/__init__.py

@@ -0,0 +1 @@
+from .replace_view_ops_with_view_copy_ops_pass import ReplaceViewOpsWithViewCopyOpsPass

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/_node_metadata_hook.py

@@ -0,0 +1,112 @@
+# mypy: allow-untyped-defs
+import contextlib
+from typing import Any
+
+import torch
+import torch.utils._pytree as pytree
+from torch._dispatch.python import enable_python_dispatcher
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.graph_module import GraphModule
+
+
+_EMPTY_NN_MODULE_STACK_KEY = "_empty_nn_module_stack_from_metadata_hook"
+
+
+def _node_metadata_hook(
+    node: torch.fx.Node,
+    metadata: dict[str, Any] | None = None,
+    fake_mode: FakeTensorMode | None = None,
+) -> None:
+    """
+    Hook for adding the appropriate metadata to nodes that are created during a
+    pass using graph.create_node. An example of how to use it:
+
+    ```
+    with _set_node_metadata_hook(gm,
+        functools.partial(_node_metadata_hook, metadata={"stack_trace": "file"})
+    ):
+        pass(gm)
+    ```
+
+    This hook should not work for all generic cases -- specifically it assumes
+    that nodes being added are only call_function nodes, and copies over the
+    first argument node's nn_module_stack.
+    """
+    # pyrefly: ignore [bad-assignment]
+    fake_mode = fake_mode or contextlib.nullcontext()
+
+    if node.op != "call_function" or not callable(node.target):
+        raise AssertionError(f"node: {node}, target: {node.target}")
+
+    if (
+        isinstance(node.target, torch._ops.OpOverload)
+        and len(node.target._schema.returns) == 0
+    ):
+        node.meta["val"] = None
+    else:
+        fake_args, fake_kwargs = pytree.tree_map_only(
+            torch.fx.Node, lambda arg: arg.meta["val"], (node.args, node.kwargs)
+        )
+        # pyrefly: ignore [bad-context-manager]
+        with fake_mode, enable_python_dispatcher():
+            fake_res = node.target(*fake_args, **fake_kwargs)
+        node.meta["val"] = fake_res
+
+    if metadata is not None:
+        for k, v in metadata.items():
+            node.meta[k] = v
+
+    # Copy over metadata from argument nodes
+    arg_meta = [
+        arg.meta
+        for arg in pytree.tree_flatten((node.args, node.kwargs))[0]
+        if isinstance(arg, torch.fx.Node)
+    ]
+    if len(arg_meta) == 0:
+        return
+    arg_meta = arg_meta[0]
+
+    node.meta["nn_module_stack"] = node.meta.get(
+        "nn_module_stack",
+        arg_meta.get(
+            "nn_module_stack",
+            {
+                _EMPTY_NN_MODULE_STACK_KEY: (
+                    _EMPTY_NN_MODULE_STACK_KEY,
+                    _EMPTY_NN_MODULE_STACK_KEY,
+                )
+            },
+        ),
+    )
+
+    node.meta["torch_fn"] = node.meta.get(
+        "torch_fn",
+        (
+            f"{node.target.__name__}_0",
+            f"{node.target.__class__.__name__}.{node.target.__name__}",
+        ),
+    )
+
+    node.meta["custom"] = node.meta.get("custom", arg_meta.get("custom", {}))
+
+
+@contextlib.contextmanager
+def _set_node_metadata_hook(gm: torch.fx.GraphModule, f):
+    """
+    Takes a callable which will be called after we create a new node. The
+    callable takes the newly created node as input and returns None.
+    """
+    if not callable(f):
+        raise AssertionError("node_metadata_hook must be a callable.")
+
+    # Add the hook to all submodules
+    for m in gm.modules():
+        if isinstance(m, GraphModule):
+            m._register_create_node_hook(f)
+    try:
+        yield
+    finally:
+        # Restore hook for all submodules
+        for m in gm.modules():
+            if isinstance(m, GraphModule):
+                m._unregister_create_node_hook(f)

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/add_runtime_assertions_for_constraints_pass.py

@@ -0,0 +1,255 @@
+# mypy: allow-untyped-defs
+import math
+import operator
+import traceback
+from functools import partial
+from typing import NamedTuple, TYPE_CHECKING
+
+import sympy
+
+import torch
+import torch.fx
+from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from torch.utils._sympy.numbers import int_oo
+from torch.utils._sympy.value_ranges import ValueRanges
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+__all__ = ["InputDim"]
+
+
+class InputDim(NamedTuple):
+    input_name: str
+    dim: int
+
+
+def _convert_to_int(val):
+    # Convert simple sympy Integers into concrete int
+    if val in (sympy.oo, int_oo):
+        return math.inf
+    if val in (-sympy.oo, -int_oo):
+        return -math.inf
+    if isinstance(val, sympy.Integer):
+        return int(val)
+    raise RuntimeError("Export constraints cannot be non-integer expressions")
+
+
+def _convert_range_to_int(range: ValueRanges):
+    if not isinstance(range, ValueRanges):
+        raise AssertionError(f"expected ValueRanges, got {type(range)}")
+    min_val = _convert_to_int(range.lower)
+    max_val = _convert_to_int(range.upper)
+    return min_val, max_val
+
+
+class _AddRuntimeAssertionsForInlineConstraintsPass(PassBase):
+    def __init__(
+        self,
+        range_constraints: dict[sympy.Symbol, ValueRanges],
+    ):
+        super().__init__()
+        self.range_constraints: dict[sympy.Symbol, ValueRanges] = range_constraints
+        self._asserts_generated_unbacked_symbols: set[sympy.Symbol] = set()
+        self.counter = 0
+
+    def _assert_range_constraint(self, node, lower, upper, assert_msg):
+        last_node = node
+        if lower > -math.inf:
+            last_node = self._insert_assert_async(
+                last_node, operator.ge, node, lower, assert_msg
+            )
+
+        if upper < math.inf:
+            last_node = self._insert_assert_async(
+                last_node, operator.le, node, upper, assert_msg
+            )
+
+    def _insert_assert_async(self, last_node, op, lower, upper, assert_msg):
+        """
+        Inserts assert_async call_function nodes in the graph. This function is
+        called **during** the interpreter-based pass.
+        """
+        self.counter += 1
+        graph = last_node.graph
+        with graph.inserting_after(last_node):
+            cmp = graph.call_function(op, (lower, upper), {})
+        with graph.inserting_after(cmp):
+            cmp_tensor = graph.call_function(
+                torch.ops.aten.scalar_tensor.default, (cmp,), {}
+            )
+        with graph.inserting_after(cmp_tensor):
+            assert_async = graph.call_function(
+                torch.ops.aten._assert_async.msg,
+                (cmp_tensor, assert_msg),
+                {},
+            )
+        return assert_async
+
+    def call(self, graph_module) -> PassResult:
+        self.existing_inline_assertions = _get_existing_inline_assertions(
+            graph_module, self.range_constraints
+        )
+
+        for module in graph_module.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if node.op != "call_function":
+                    continue
+                if "val" not in node.meta:
+                    continue
+
+                val = node.meta["val"]
+                # In general, we may have to deal the case such as: ret[1].shape[0].
+                # We need first find out what symbols require assertion, then we need to follow the path
+                # from ret to the symbol, construct the proxies along the way and construct the messages
+                # piece-wise at the same time.
+                #
+                # We use post-order traversal to collect all the proxies callbacks needed, construct
+                # the error message callbacks, and at the top-level traversal tree we execute all the callbacks.
+                # We need the callbacks because, in order to call the function to create a proxy for shape[0], we
+                # need the proxy for shape, which further requires the proxy for ret[1], etc.
+
+                def add_assertions(val):
+                    call_backs: list[Callable] = []
+                    messages: list[str] = []
+                    if isinstance(val, (torch.SymInt, torch.SymFloat, torch.SymBool)):
+                        symbol = val.node.expr
+                        if symbol in self.existing_inline_assertions:
+                            return call_backs, messages
+                        if isinstance(symbol, sympy.Symbol) and free_unbacked_symbols(
+                            symbol
+                        ):
+                            if symbol in self._asserts_generated_unbacked_symbols:
+                                return call_backs, messages
+                            # We only care about unbacked symints for these inline
+                            # constraints, which are prefixed with 'u'
+                            constraint = self.range_constraints[symbol]
+                            min_val, max_val = _convert_range_to_int(constraint)
+                            assert_msg = f" is outside of inline constraint [{min_val}, {max_val}]."
+                            call_backs.append(
+                                partial(
+                                    self._assert_range_constraint,
+                                    lower=min_val,
+                                    upper=max_val,
+                                )
+                            )
+                            messages.append(assert_msg)
+                            self._asserts_generated_unbacked_symbols.add(symbol)
+
+                    elif isinstance(val, torch.Tensor):
+                        for i, sym in enumerate(val.shape):
+                            cbs, msgs = add_assertions(sym)
+                            for cb, msg in zip(cbs, msgs):
+
+                                def sym_size_cb(node, assert_msg, dim):
+                                    with node.graph.inserting_after(node):
+                                        dim_node = module.graph.call_function(
+                                            torch.ops.aten.sym_size.int,
+                                            (node, dim),
+                                            {},
+                                        )
+                                    cb(node=dim_node, assert_msg=assert_msg)
+
+                                call_backs.append(partial(sym_size_cb, dim=i))
+                                messages.append(f".shape[{i}]" + msg)
+                    return call_backs, messages
+
+                callbacks, messages = add_assertions(val)
+                for cb, msg in zip(callbacks, messages):
+                    cb(node=node, assert_msg=f"{node}" + msg)
+
+            module.recompile()
+
+        # Sometimes this pass would return a wrong graph where we have mismatched
+        # node names in signature. Before we fix it, let's just skip it.
+        if (
+            self.counter == 0
+            and type(self) is _AddRuntimeAssertionsForInlineConstraintsPass
+        ):
+            return PassResult(graph_module, False)
+
+        # Populate the stack trace with dummy vals to respect IR
+        for node in graph_module.graph.nodes:
+            if not node.meta.get("stack_trace", None) and node.op not in [
+                "placeholder",
+                "output",
+            ]:
+                node.meta["stack_trace"] = "".join(traceback.format_stack(limit=1))
+        return PassResult(graph_module, True)
+
+
+def _get_existing_inline_assertions(
+    graph_module: torch.fx.GraphModule,
+    range_constraints: dict[sympy.Symbol, ValueRanges],
+) -> dict[sympy.Symbol, ValueRanges]:
+    existing_inline_assertions: dict[sympy.Symbol, ValueRanges] = {}
+
+    for module in graph_module.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            continue
+
+        # Find all the existing inline assertions. They will look something like:
+        # %_local_scalar_dense = call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%arg1_1,), kwargs = {})
+        # %ge = call_function[target=operator.ge](args = (%_local_scalar_dense, 0), kwargs = {})
+        # %_assert_scalar = call_function[target=torch.ops.aten._assert_scalar.default](args = (%scalar_tensor, "..."), kwargs = {})
+        for node in module.graph.nodes:
+            if node.target != torch.ops.aten._assert_scalar.default:
+                continue
+
+            compare_arg = node.args[0]
+            if not (
+                isinstance(compare_arg, torch.fx.Node)
+                and compare_arg.op == "call_function"
+                and compare_arg.target in (operator.le, operator.ge)
+                and len(compare_arg.args) == 2
+            ):
+                continue
+
+            compare_op = compare_arg.target
+            lhs, rhs = compare_arg.args
+
+            def maybe_get_symint(x):
+                if (
+                    isinstance(x, torch.fx.Node)
+                    and "val" in x.meta
+                    and isinstance(x.meta["val"], torch.SymInt)
+                ):
+                    return x.meta["val"].node.expr
+                return x
+
+            lhs = maybe_get_symint(lhs)
+            rhs = maybe_get_symint(rhs)
+
+            if compare_op is operator.ge:
+                lhs, rhs = rhs, lhs
+
+            if isinstance(lhs, sympy.Symbol) and isinstance(rhs, int):
+                symint = lhs
+                scalar = rhs
+            elif isinstance(rhs, sympy.Symbol) and isinstance(lhs, int):
+                symint = rhs
+                scalar = lhs
+            else:
+                continue
+
+            if symint not in range_constraints:
+                raise RuntimeError(
+                    f"Unable to find symint {symint} in {range_constraints}"
+                )
+
+            previous_range = existing_inline_assertions.get(
+                symint, ValueRanges(-math.inf, math.inf)
+            )
+
+            if symint is lhs:
+                bounds = ValueRanges(-math.inf, scalar)
+            else:
+                bounds = ValueRanges(scalar, math.inf)
+            existing_inline_assertions[symint] = previous_range & bounds
+
+    return existing_inline_assertions

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/collect_tracepoints_pass.py

@@ -0,0 +1,158 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import operator
+from typing import TYPE_CHECKING
+
+import torch
+from torch.export.exported_program import ConstantArgument, TensorArgument
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+
+
+if TYPE_CHECKING:
+    from torch.export.exported_program import ModuleCallSignature
+    from torch.export.graph_signature import ExportGraphSignature
+
+
+__all__ = ["CollectTracepointsPass"]
+
+
+class CollectTracepointsPass(PassBase):
+    """
+    Performs constant folding and constant propagation.
+    """
+
+    def __init__(
+        self, specs: dict[str, ModuleCallSignature], sig: ExportGraphSignature
+    ) -> None:
+        super().__init__()
+        self.specs = specs
+        self.sig = sig
+
+    def call(self, gm: torch.fx.GraphModule) -> PassResult | None:
+        def get_arg_spec(arg) -> TensorArgument | ConstantArgument:
+            if isinstance(arg, torch.fx.Node):
+                if isinstance(arg.meta.get("val"), torch.Tensor):
+                    return TensorArgument(name=arg.name)
+                else:
+                    raise AssertionError(
+                        "Symint input is not implemented yet for submodule call signature."
+                    )
+            else:
+                return ConstantArgument(name="", value=arg)
+
+        for module in gm.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            nn_module_stack = None
+            for node in module.graph.nodes:
+                if node.op != "call_function":
+                    continue
+                if node.target is torch.ops.higher_order._export_tracepoint:
+                    kind = node.kwargs["kind"]
+                    if kind == "module_call_outputs":
+                        nn_module_stack = node.meta["nn_module_stack"]
+                    elif kind == "module_call_inputs":
+                        nn_module_stack = None
+                    else:
+                        raise AssertionError(f"Unknown tracepoint kind: {kind}")
+                elif node.meta["nn_module_stack"] == nn_module_stack:
+                    node.meta["nn_module_stack"].popitem()
+                else:
+                    nn_module_stack = None
+            nn_module_stack = None
+            for node in reversed(module.graph.nodes):
+                if node.op != "call_function":
+                    continue
+                if node.target is torch.ops.higher_order._export_tracepoint:
+                    kind = node.kwargs["kind"]
+                    if kind == "module_call_inputs":
+                        nn_module_stack = node.meta["nn_module_stack"]
+                    elif kind == "module_call_outputs":
+                        nn_module_stack = None
+                    else:
+                        raise AssertionError(f"Unknown tracepoint kind: {kind}")
+                elif node.meta["nn_module_stack"] == nn_module_stack:
+                    node.meta["nn_module_stack"].popitem()
+                else:
+                    nn_module_stack = None
+
+        def copy_sig(sig) -> ModuleCallSignature:
+            from torch.export.exported_program import ModuleCallSignature
+
+            return ModuleCallSignature(
+                inputs=[],
+                outputs=[],
+                in_spec=sig.in_spec,
+                out_spec=sig.out_spec,
+                forward_arg_names=None,
+            )
+
+        for module in gm.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if node.op != "call_function":
+                    continue
+                if node.target is torch.ops.higher_order._export_tracepoint:
+                    # There's some subtlety worth noting. Here fqn corresponds to
+                    # the call name, whereas path corresponds to the module name.
+                    # They are not necessarily the same! When a submodule is shared
+                    # through different aliases, there are as many _export_tracepoint
+                    # markers as there are aliases, since the shared submodule is
+                    # wrapped once for each alias.
+                    path = node.kwargs["path"]
+                    fqn, _ = next(reversed(node.meta["nn_module_stack"].values()))
+
+                    module_key = next(reversed(node.meta["nn_module_stack"]))
+                    if "@" in module_key:
+                        suffix = module_key.split("@")[-1]
+                        path = f"{path}@{suffix}"
+
+                        call_fqn = f"{fqn}@{suffix}"
+                        if call_fqn not in self.specs:
+                            self.specs[call_fqn] = copy_sig(self.specs[fqn])
+                        fqn = call_fqn
+
+                    kind = node.kwargs["kind"]
+                    for i, arg in enumerate(node.args):
+                        # We only update the signature of the alias used to call
+                        # the submodule. Otherwise the signatures of all aliases
+                        # would get conflated; the inputs/outputs of every call
+                        # would be recorded in every other call as well.
+                        if fqn == path:
+                            if kind == "module_call_inputs":
+                                self.specs[path].inputs.append(get_arg_spec(arg))
+                            elif kind == "module_call_outputs":
+                                self.specs[path].outputs.append(get_arg_spec(arg))
+                            else:
+                                raise AssertionError(f"Unknown tracepoint kind: {kind}")
+                        if isinstance(arg, torch.fx.Node):
+                            for user in node.users:
+                                if user.op != "call_function":
+                                    raise AssertionError(
+                                        f"expected call_function, got {user.op}"
+                                    )
+                                if user.target is not operator.getitem:
+                                    raise AssertionError(
+                                        f"expected getitem target, got {user.target}"
+                                    )
+                                if not isinstance(user.args[1], int):
+                                    raise AssertionError(
+                                        f"expected int arg, got {type(user.args[1])}"
+                                    )
+                                if user.args[1] == i:
+                                    user.replace_all_uses_with(arg)
+                                    self.sig.replace_all_uses(user.name, arg.name)
+                                    break
+                    users = list(node.users)
+                    for user in users:
+                        if len(user.users) != 0:
+                            raise AssertionError(
+                                f"expected no users, got {len(user.users)}"
+                            )
+                        gm.graph.erase_node(user)
+                    gm.graph.erase_node(node)
+            return PassResult(gm, True)
+
+        return None

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/constant_folding.py

@@ -0,0 +1,304 @@
+# mypy: allow-untyped-defs
+import collections
+from collections import defaultdict
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.utils._pytree as pytree
+
+
+aten = torch.ops.aten
+
+# We would like to split modules into two subgraphs for runtime weight updates to work correctly.
+# The use case and more information could be found at:
+# https://docs.google.com/document/d/1inZC-8KarJ6gKB7G9egmYLx1V_dKX_apxon0w4zPC0Q/edit?usp=sharing
+META_TAG = "MODULE_TYPE"
+MODULE_TAG = "_MAIN_MODULE"
+CONST_MODULE_TAG = "_CONST_MODULE"
+
+
+def replace_node_with_constant(gm, node, constant, name=None):
+    g = gm.graph
+
+    if name:
+        qualname = name
+    else:
+        if not hasattr(gm, "_frozen_param_count"):
+            gm._frozen_param_count = 0
+        i = gm._frozen_param_count
+
+        while True:
+            qualname = f"_frozen_param{i}"
+            if not hasattr(gm, qualname):
+                break
+            i += 1
+
+        gm._frozen_param_count = i + 1
+
+    with g.inserting_before(node):
+        new_input_node = g.create_node("get_attr", qualname, (), {})
+        node.replace_all_uses_with(new_input_node)
+        new_input_node.meta.update(node.meta)
+        g.erase_node(node)
+
+    # needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning
+    gm.register_buffer(qualname, constant)
+    setattr(gm, qualname, constant)
+
+
+class ConstantFolder(torch.fx.Interpreter):
+    def __init__(
+        self,
+        gm: torch.fx.GraphModule,
+        skip_constructors: bool = False,
+    ):
+        super().__init__(gm)
+        self.node_replacements: dict[torch.fx.Node, Any] = {}
+        self.replaced_uses: dict[torch.fx.Node, int] = collections.Counter()
+        self.unknown_value = object()
+        self.skip_constructors: bool = skip_constructors
+
+        # overwrite this to deallocate env values if their only remaining use
+        # is the output
+        self.user_to_last_uses = self.node_to_last_non_output_use()
+
+    def is_impure(self, node: torch.fx.Node) -> bool:
+        if (
+            node.target is torch.ops.prims.convert_element_type.default
+            and node.args[0].op == "get_attr"  # type: ignore[union-attr]
+            and node.args[0].meta["val"].dtype == torch.int8  # type: ignore[union-attr]
+            and node.args[1] == torch.bfloat16
+        ):
+            # For int8_weight -> dq -> bf16_weight
+            return True
+        if node.target in [
+            torch.ops.quantized_decomposed.dequantize_per_channel.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.tensor,
+            torch.ops.pt2e_quant.dequantize_affine,
+        ]:
+            # For the pattern fp32_weight -> q -> dq
+            # We only folding fp32_weight -> q
+            # int8_weight and leave dq in graph to be fused
+            return True
+        return False
+
+    def node_to_last_non_output_use(self):
+        last_non_output_use = collections.defaultdict(list)
+        seen_uses = set()
+        output_node = next(iter(reversed(self.module.graph.nodes)))  # type: ignore[arg-type, union-attr]
+
+        for node in reversed(self.module.graph.nodes):  # type: ignore[arg-type, union-attr]
+            if node.target == "output":
+                continue
+
+            def add_use(inp):
+                if inp in seen_uses:
+                    return
+
+                seen_uses.add(inp)
+                last_non_output_use[node].append(inp)
+
+            # In-place is fine since we don't mutate
+            pytree.tree_map_only_(torch.fx.Node, add_use, (node.args, node.kwargs))
+
+            # if this node is only used in output, we want to gc it right away
+            if len(node.users) == 1 and output_node in node.users:
+                last_non_output_use[node].append(node)
+
+        return last_non_output_use
+
+    def run_node(self, node):
+        if node.target == "output":
+            # because we remove nodes from env on last non output use,
+            # re-define them now or we'll get error in interpreter
+            def set_env(arg):
+                self.env[arg] = self.unknown_value
+
+            # In-place is fine since we don't mutate
+            pytree.tree_map_only_(torch.fx.Node, set_env, node.args)
+            return super().run_node(node)
+
+        args, kwargs = self.fetch_args_kwargs_from_env(node)
+        flattened_inputs = pytree.arg_tree_leaves(*args, **kwargs)
+
+        # We need to do this weird thing because in cases where flattened_inputs
+        # contains a ScriptObject, equality checking results in a type error if
+        # the types are different.
+        if any(
+            type(self.unknown_value) is type(input_) and self.unknown_value == input_
+            for input_ in flattened_inputs
+        ):
+            return self.unknown_value
+
+        # TODO - fix errors with this
+        if (
+            node.op == "call_function"
+            and node.target is aten._efficientzerotensor.default
+        ):
+            return self.unknown_value
+
+        # TODO - constant folding triton kernel returns the inputs -- fix this
+        if (
+            node.op == "call_function"
+            and node.name == "triton_kernel_wrapper_functional_proxy"
+        ):
+            return self.unknown_value
+
+        # skip constructors, since inductor generates optimal code for them already
+        # and turning into tensor would result in an additional global memory read
+        # TODO - more complicated strategy
+        if (
+            self.skip_constructors
+            and node.op != "get_attr"
+            and not any(isinstance(e, torch.Tensor) for e in flattened_inputs)
+        ):
+            return self.unknown_value
+
+        # All mutations should either be removed or on inputs which we did not make constant
+        if (
+            isinstance(node.target, torch._ops.OpOverload)
+            and torch.Tag.nondeterministic_seeded in node.target.tags
+        ):
+            return self.unknown_value
+
+        out = super().run_node(node)
+
+        if node.op != "get_attr" and isinstance(out, torch.Tensor):
+            if out.device.type == "meta":
+                return out
+
+            if not self.insertable_tensor_check(out):
+                return out
+
+            if self.is_impure(node):
+                return self.unknown_value
+
+            self.add_node_replacement(node, out)
+
+            flattened_node_inps = pytree.arg_tree_leaves(*node.args, **node.kwargs)
+
+            for n in flattened_node_inps:
+                if not isinstance(n, torch.fx.Node):
+                    continue
+
+                self.replaced_uses[n] += 1
+
+            for to_delete in self.user_to_last_uses.get(node, []):
+                if self.replaced_uses[to_delete] == len(to_delete.users):
+                    self.node_replacements.pop(to_delete, None)
+
+        return out
+
+    def insertable_tensor_check(self, tensor: torch.Tensor) -> bool:
+        return True
+
+    def add_node_replacement(self, node: torch.fx.Node, tensor: torch.Tensor) -> None:
+        self.node_replacements[node] = tensor
+
+    def run(self):  # type: ignore[override]
+        env = {}
+        for n in self.module.graph.find_nodes(op="placeholder"):  # type: ignore[operator, union-attr]
+            env[n] = self.unknown_value
+        return super().run(initial_env=env)
+
+
+def constant_fold(
+    gm: torch.fx.GraphModule,
+    constraint_fn: Callable[[torch.fx.Node], bool] | None = None,
+):
+    with torch.utils._python_dispatch._disable_current_modes():
+        cf = ConstantFolder(gm, skip_constructors=True)
+        cf.run()
+
+        for node, constant in cf.node_replacements.items():
+            if constraint_fn is not None and not constraint_fn(node):
+                continue
+            replace_node_with_constant(gm, node, constant)
+
+        erased_params = []
+        # Get all attr users by looking up the graph instead from node.users, because in this case
+        # _tensor_constant0 and _tensor_constant0_1 are actually refereing to the same tensor.
+
+        #     opcode         name                 target            args                         kwargs
+        # -------------  -------------------  ----------------  ---------------------------  --------
+        # placeholder    arg0_1               arg0              ()                           {}
+        # get_attr       _tensor_constant0    state             ()                           {}
+        # call_function  add                  aten.add.Tensor   (arg0_1, _tensor_constant0)  {}
+        # get_attr       _tensor_constant0_1  state             ()                           {}
+        # call_function  add_                 aten.add_.Tensor  (_tensor_constant0_1, 1)     {}
+        # output         output               output            ([add],)                     {}
+
+        get_attr_node_users = defaultdict(list)
+        for node in gm.graph.nodes:
+            if node.op == "get_attr":
+                get_attr_node_users[node.target].extend(node.users.keys())
+        for node in gm.graph.find_nodes(op="get_attr"):
+            if node.op == "get_attr" and len(get_attr_node_users[node.target]) == 0:
+                if hasattr(gm, node.target):
+                    delattr(gm, node.target)
+                erased_params.append(node)
+        for node in erased_params:
+            gm.graph.erase_node(node)
+
+        gm.graph.eliminate_dead_code()
+        gm.graph.lint()
+        gm.recompile()
+
+
+def constant_graph_tag(gm: torch.fx.GraphModule) -> None:
+    with torch.utils._python_dispatch._disable_current_modes():
+        cf = ConstantFolder(gm, skip_constructors=True)
+        cf.run()
+
+        for node in gm.graph.nodes:
+            if (
+                node.op == "get_attr"
+                or node in cf.node_replacements
+                or node in cf.replaced_uses
+            ):
+                node.meta[META_TAG] = CONST_MODULE_TAG
+            else:
+                node.meta[META_TAG] = MODULE_TAG
+
+
+def run_and_get_constant_graph(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    """
+    Construct a GraphModule which corresponds to the part which could be
+    constant folded in provided gm.
+    """
+
+    constant_graph_tag(gm)
+    # We rewrite the tags, if it's a constant being directly consumed, without
+    # any folding opportunity, we keep it in main gm.
+    for node in gm.graph.find_nodes(op="get_attr"):
+        used_to_fold = False
+        for u in node.users:
+            if u.meta[META_TAG] == CONST_MODULE_TAG:
+                used_to_fold = True
+                break
+        if not used_to_fold:
+            node.meta[META_TAG] = MODULE_TAG
+
+    new_graph = torch.fx.Graph()
+
+    node_remapping: dict[torch.fx.Node, torch.fx.Node] = {}
+    output_nodes = []
+    for node in gm.graph.nodes:
+        if node.meta[META_TAG] == MODULE_TAG:
+            continue
+
+        new_node = new_graph.node_copy(node, lambda x: node_remapping[x])
+        node_remapping[node] = new_node
+
+        for user in node.users:
+            if user.meta[META_TAG] == MODULE_TAG:
+                output_nodes.append(new_node)
+                break
+
+    new_graph.output(tuple(output_nodes))
+    new_graph.lint()
+    new_gm = torch.fx.GraphModule(gm, new_graph)
+
+    return new_gm

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/functionalize_side_effectful_ops_pass.py

@@ -0,0 +1,100 @@
+import copy
+
+import torch
+from torch._export.pass_base import (
+    _ExportPassBaseDeprecatedDoNotUse,
+    Argument,
+    PassResult,
+)
+from torch._export.pass_infra.node_metadata import NodeMetadata
+from torch._export.pass_infra.proxy_value import ProxyValue
+from torch._ops import OpOverload
+
+
+aten = torch.ops.aten
+
+_NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS: dict[OpOverload, OpOverload] = {
+    aten.sym_constrain_range.default: aten._functional_sym_constrain_range.default,
+    aten._assert_async.msg: aten._functional_assert_async.msg,
+}
+
+
+class _FunctionalizeSideEffectfulOpsPass(_ExportPassBaseDeprecatedDoNotUse):
+    """
+    Functionalize ops with side effect in graph module by replacing the op with
+    functional version of it. A new dependency token (`dep_token`) will be
+    created and propagated through functional ops to output.
+    For example:
+    ```
+    def f(x):
+        sym_constrain_range(x.shape[0], min=1, max=3)
+        return x.add(3)
+    ```
+    Will be transformed to:
+    ```
+    def f(x):
+        dep_token0 = _make_dep_token()
+        dep_token1 = _functional_sym_constrain_range(
+            x.shape[0], min=1, max=3, dep_token=dep_token0
+        )
+
+        return x.add(3), dep_token1
+    ```
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._dep_token: ProxyValue | None = None
+        self._next_dep_token_index: int | None = None
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        # Early return if no non-functional assertions.
+        if not any(
+            n.target in _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS
+            for n in graph_module.graph.nodes
+        ):
+            return PassResult(graph_module=graph_module, modified=False)
+
+        gm = copy.deepcopy(graph_module)
+        self._dep_token = None
+        self._next_dep_token_index = None
+        return super().call(gm)
+
+    def call_operator(
+        self,
+        op: OpOverload,
+        args: tuple[Argument, ...],
+        kwargs: dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        if op not in _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS:
+            return super().call_operator(op, args, kwargs, meta)
+
+        if self._dep_token is None:
+            self._dep_token = super().call_operator(
+                aten._make_dep_token,
+                args=(),
+                kwargs={},
+                meta=self._create_dummy_node_metadata(),
+            )
+            self._dep_token.node.name = "dep_token0"
+            self._next_dep_token_index = 1
+
+        self._dep_token = super().call_operator(
+            _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS[op],
+            args=args,
+            kwargs={**kwargs, "dep_token": self._dep_token},
+            meta=meta,
+        )
+        if self._next_dep_token_index is None:
+            raise AssertionError("_next_dep_token_index must not be None")
+        self._dep_token.node.name = f"dep_token{self._next_dep_token_index}"
+        self._next_dep_token_index += 1
+
+        return self._dep_token
+
+    def output(self, results: list[Argument], meta: NodeMetadata) -> ProxyValue:
+        if self._dep_token is None:
+            raise AssertionError("_dep_token must not be None")
+
+        return super().output(results=(*results, self._dep_token), meta=meta)  # type: ignore[arg-type]

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/insert_custom_op_guards.py

@@ -0,0 +1,84 @@
+import functools
+from collections import defaultdict
+
+import torch
+from torch._export.passes._node_metadata_hook import (
+    _node_metadata_hook,
+    _set_node_metadata_hook,
+)
+from torch._library.fake_profile import OpProfile, TensorMetadata
+
+
+def insert_custom_op_guards(gm: torch.fx.GraphModule, ops_to_guard: set[str]) -> None:
+    """
+    This is used by draft_export to insert guards in front of calls to custom
+    operators which have a generated fake kernel.
+    """
+    for node in gm.graph.nodes:
+        if node.op == "call_function" and str(node.target) in ops_to_guard:
+            with (
+                _set_node_metadata_hook(
+                    gm,
+                    functools.partial(
+                        _node_metadata_hook,
+                        metadata={"stack_trace": node.meta.get("stack_trace")},
+                    ),
+                ),
+                gm.graph.inserting_before(node),
+            ):
+                for arg in (*node.args, *node.kwargs.values()):
+                    if isinstance(arg, torch.fx.Node) and isinstance(
+                        arg.meta.get("val"), torch.Tensor
+                    ):
+                        val = arg.meta["val"]
+                        gm.graph.call_function(
+                            torch.ops.aten._assert_tensor_metadata.default,
+                            args=(arg,),
+                            kwargs={
+                                "dtype": val.dtype,
+                                "device": val.device,
+                                "layout": val.layout,
+                            },
+                        )
+
+    gm.recompile()
+
+
+def get_op_profiles(
+    gm: torch.fx.GraphModule, ops_to_guard: set[str]
+) -> dict[str, set[OpProfile]]:
+    """
+    This is used by draft_export to get a list of custom operator profiles so
+    that we can generate fake kernels.
+    """
+
+    def _get_op_profile(node: torch.fx.Node) -> OpProfile:
+        args_profile = tuple(
+            TensorMetadata.maybe_from_tensor(arg.meta.get("val"))
+            if isinstance(arg, torch.fx.Node)
+            else None
+            for arg in (*node.args, *node.kwargs.values())
+        )
+
+        out_profile = None
+        meta = node.meta.get("val")
+        if meta is None:
+            raise AssertionError("node.meta['val'] must not be None")
+        if isinstance(meta, torch.Tensor):
+            out_profile = TensorMetadata.maybe_from_tensor(meta)
+        elif isinstance(meta, (list, tuple)):
+            out_profile = tuple(TensorMetadata.maybe_from_tensor(m) for m in meta)  # type: ignore[assignment]
+        if out_profile is None:
+            raise AssertionError(
+                f"out_profile must not be None for meta type {type(meta)}"
+            )
+
+        return OpProfile(args_profile, out_profile)  # type: ignore[arg-type]
+
+    op_profiles: dict[str, set[OpProfile]] = defaultdict(set)
+
+    for node in gm.graph.nodes:
+        if node.op == "call_function" and str(node.target) in ops_to_guard:
+            op_profiles[str(node.target)].add(_get_op_profile(node))
+
+    return op_profiles

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/lift_constants_pass.py

@@ -0,0 +1,449 @@
+# mypy: allow-untyped-defs
+import collections
+import logging
+from typing import Any
+
+import torch
+from torch._export.verifier import SpecViolationError
+from torch._guards import detect_fake_mode
+from torch._library.fake_class_registry import FakeScriptObject, maybe_to_fake_obj
+from torch._library.opaque_object import (
+    get_opaque_type_name,
+    is_opaque_reference_type,
+    is_opaque_type,
+)
+from torch._subclasses.fake_tensor import unset_fake_temporarily
+from torch.export.exported_program import (
+    ArgumentSpec,
+    CustomObjArgument,
+    ExportGraphSignature,
+    InputKind,
+    InputSpec,
+    TensorArgument,
+)
+from torch.fx._symbolic_trace import _ConstantAttributeType
+from torch.fx.graph_module import _get_attr
+
+
+log = logging.getLogger(__name__)
+
+
+class ConstantAttrMap(collections.abc.MutableMapping):
+    """A mapping class that understands how to use module constants (tensors,
+    ScriptObjects, FakeScriptObjects, opaque objects) as keys. We store tensors,
+    FakeScriptObjects, and opaque objects normally, but ScriptObjects are stored
+    by hash, because different torch.ScriptObjects can point to the same
+    underlying value (but we guarantee that they will `hash()` to the same value
+    if that's the case).
+    """
+
+    def __init__(self) -> None:
+        # Underlying dict that we use to implement this mapping.
+        self._constant_attrs: dict[
+            int | torch.Tensor | FakeScriptObject | torch.utils._pytree.TreeSpec,
+            list[Any],
+        ] = {}
+        # Map from the hash(ScriptObject) to the ScriptObject itself. Used for
+        # APIs like `__iter__` that should look like they're returning the
+        # original ScriptObjects.
+        self._script_object_map: dict[int, torch.ScriptObject] = {}
+
+    def __getitem__(self, key: _ConstantAttributeType) -> Any:
+        real_key = hash(key) if isinstance(key, torch.ScriptObject) else key
+        if not isinstance(
+            real_key, (int, torch.Tensor, FakeScriptObject)
+        ) and not is_opaque_type(type(real_key)):
+            raise AssertionError(
+                f"expected int, Tensor, FakeScriptObject, or opaque type key, got {type(real_key)}"
+            )
+        return self._constant_attrs[real_key]
+
+    def __setitem__(self, key: _ConstantAttributeType, value):
+        # we shouldn't actually call this, should go to add() instead to handle aliasing
+        raise NotImplementedError(
+            """Directly setting values for ConstantAttrMap is not supported, please use add(key, value) instead.
+The same key can be mapped to multiple values, for handling constant aliasing."""
+        )
+
+    def add(self, key: _ConstantAttributeType, value: Any) -> None:
+        if isinstance(key, torch.ScriptObject):
+            if hash(key) not in self._constant_attrs:
+                self._constant_attrs[hash(key)] = []
+            self._constant_attrs[hash(key)].append(value)
+            self._script_object_map[hash(key)] = key
+        elif isinstance(key, (torch.Tensor, FakeScriptObject)) or is_opaque_type(
+            type(key)
+        ):
+            if key not in self._constant_attrs:
+                self._constant_attrs[key] = []
+            self._constant_attrs[key].append(value)
+        else:
+            raise TypeError(
+                f"Expected key to be a tensor or ScriptObject, got {type(key)}"
+            )
+
+    def __delitem__(self, key: _ConstantAttributeType):
+        real_key = hash(key) if isinstance(key, torch.ScriptObject) else key
+
+        del self._constant_attrs[real_key]
+
+    def __iter__(self):
+        for key in self._constant_attrs:
+            if isinstance(key, int):
+                yield self._script_object_map[key]
+            else:
+                yield key
+
+    def __len__(self):
+        return len(self._constant_attrs)
+
+    def __contains__(self, key: object) -> bool:
+        real_key = hash(key) if isinstance(key, torch.ScriptObject) else key
+        return real_key in self._constant_attrs
+
+
+def get_constant_fqn(node: torch.fx.Node, constant_name: str) -> str:
+    # The FQN of the constant tensor in the state dict should
+    # correspond to the module where the constant tensor was
+    # originally used.
+    if len(node.meta["nn_module_stack"]) == 0:
+        return constant_name
+    parent_fqn = list(node.meta["nn_module_stack"].values())[-1][0]
+    if len(parent_fqn) > 0:
+        return f"{parent_fqn}.{constant_name}"
+    else:
+        return constant_name
+
+
+def _get_first_fqn(
+    const_attrs: ConstantAttrMap,
+    key: _ConstantAttributeType,
+) -> Any:
+    fqns = const_attrs.get(key)
+    return fqns[0] if fqns else None
+
+
+def _unused_constant(node: torch.fx.Node) -> list[torch.fx.Node] | None:
+    """
+    If there is a tensor constant created while tracing, here is how the graph
+    looks like:
+
+        %_tensor_constant0 : [num_users=1] = get_attr[target=_tensor_constant0]
+        %lift_fresh_copy : [num_users=1] = call_function[target=torch.ops.aten.lift_fresh_copy.default](args = (%_tensor_constant0,))
+        %detach_ : [num_users=?] = call_function[target=torch.ops.aten.detach_.default](args = (%lift_fresh_copy,))
+
+    To check to see if the tensor constant is being used, we want to traverse to
+    the detach node to see if it's actually being used.
+
+    This function returns None if this constant is being used, otherwise it returns the
+    lift_fresh and detach node to be removed later.
+    """  # noqa: B950
+    if len(node.users) > 1:
+        return None
+
+    lift_fresh_node = next(iter(node.users.keys()))
+    if not (
+        lift_fresh_node.op == "call_function"
+        and lift_fresh_node.target
+        in (
+            torch.ops.aten.lift_fresh.default,
+            torch.ops.aten.lift_fresh_copy.default,
+        )
+    ):
+        return None
+
+    if len(lift_fresh_node.users) > 1:
+        return None
+
+    # Case 1: lift node is not used anywhere
+    if len(lift_fresh_node.users) == 0:
+        return [lift_fresh_node, node]
+
+    detach_node = next(iter(lift_fresh_node.users.keys()))
+    if not (
+        detach_node.op == "call_function"
+        and detach_node.target
+        in (
+            torch.ops.aten.detach_.default,
+            torch.ops.aten.detach.default,
+        )
+    ):
+        return None
+
+    if len(detach_node.users) > 0:
+        return None
+    else:
+        # Case 2: Lift node's child is not used anywhere
+        return [detach_node, lift_fresh_node, node]
+
+
+def lift_constants_pass(
+    gm: torch.fx.GraphModule,
+    graph_signature: ExportGraphSignature,
+    constant_attrs: ConstantAttrMap,
+) -> dict[str, _ConstantAttributeType]:
+    """
+    Takes a graph module, graph signature, and modifies them inplace to lift any
+    constants (tensors or custom classes) as inputs to the graph. Returns a
+    dictionary of names to constants.
+
+    Arguments:
+        gm (torch.fx.GraphModule): The graph module containing the graph and constants to lift.
+        graph_signature (ExportGraphSignature): This graph signature will be
+            mutated to add additional CONSTANT_TENSOR and CUSTOM_OBJ inputs.
+        constant_attrs (ConstantAttr): A mapping from a constant value to its
+            fully-qualified path in `gm`. This is used to maintain consistent
+            location of constants between the original module and the exported
+            version.
+
+    Returns:
+        A dictionary of fqn => constant value.
+    """
+    all_constants: dict[str, _ConstantAttributeType] = {}
+
+    input_specs = graph_signature.input_specs
+    num_custom_obj = sum(
+        input_spec.kind == InputKind.CUSTOM_OBJ for input_spec in input_specs
+    )
+    num_tensor_constants = sum(
+        input_spec.kind == InputKind.CONSTANT_TENSOR for input_spec in input_specs
+    )
+
+    fake_mode = detect_fake_mode(
+        tuple(node.meta["val"] for node in gm.graph.nodes if node.op == "placeholder")
+    )
+
+    first_user_input_loc, first_user_input = 0, next(iter(gm.graph.nodes))
+    used_target_names = set()
+
+    input_nodes = [node for node in gm.graph.nodes if node.op == "placeholder"]
+    if len(input_nodes) != len(input_specs):
+        raise AssertionError(
+            f"input nodes count {len(input_nodes)} != input specs count {len(input_specs)}"
+        )
+    for i, (node, input_spec) in enumerate(zip(input_nodes, input_specs)):
+        used_target_names.add(input_spec.target)
+        if input_spec.kind == InputKind.USER_INPUT:
+            first_user_input = node
+            first_user_input_loc = i
+            break
+
+    lifted_objs = ConstantAttrMap()
+    renamed_targets = {}
+    for node in list(gm.graph.nodes):
+        if node.op == "get_attr":
+            if nodes_to_remove := _unused_constant(node):
+                # Remove the node if it's not being used
+                for node_rm in nodes_to_remove:
+                    gm.graph.erase_node(node_rm)
+                continue
+
+            constant_val = _get_attr(gm, node.target)
+            # These are not hashable and not gonna be lifted
+            # so we can skip them earlier
+            if isinstance(constant_val, torch.fx.GraphModule):
+                continue
+            if "LoweredBackendModule" in type(constant_val).__name__:
+                continue
+            if "AOTInductorRunnerWrapper" in type(constant_val).__name__:
+                continue
+            if isinstance(constant_val, torch.utils._pytree.TreeSpec):
+                continue
+
+            if constant_val in lifted_objs:
+                # We already lifted this constant elsewhere. Just rewrite uses
+                # of this get_attr to point to the already-existing placeholder
+                # node.
+                const_placeholder_node = _get_first_fqn(lifted_objs, constant_val)
+                node.replace_all_uses_with(const_placeholder_node)
+                gm.graph.erase_node(node)
+                renamed_targets[node.name] = const_placeholder_node.name
+                continue
+
+            # For ScriptObject, Tensor and FakeScriptObject constants:
+            # First check if the constant was an attribute on some module by
+            # consulting `constant_attrs` map. If it is, use the fqn that keeps
+            # its location consistent with the eager module.
+            #
+            # If it's not in the `constant_attrs` map, that means it's an inline
+            # constant (e.g. x + torch.tensor(0)), and thus did not have a
+            # specific location in the eager module. In that case, just generate
+            # some name and attach it to the module in which it was used.
+            if isinstance(
+                constant_val, (torch.ScriptObject, FakeScriptObject)
+            ) or is_opaque_reference_type(type(constant_val)):
+                constant_kind = InputKind.CUSTOM_OBJ
+                constant_fqn = _get_first_fqn(constant_attrs, constant_val)
+                if constant_fqn is not None:
+                    constant_name = constant_fqn.replace(".", "_")
+                else:
+                    constant_name = f"lifted_custom_{num_custom_obj}"
+                    constant_fqn = get_constant_fqn(node, constant_name)
+                    while constant_fqn in used_target_names:
+                        num_custom_obj += 1
+                        constant_name = f"lifted_custom_{num_custom_obj}"
+                        constant_fqn = get_constant_fqn(node, constant_name)
+                    num_custom_obj += 1
+            elif isinstance(constant_val, torch.Tensor):
+                # Remove the parameterness of constant_val
+                if isinstance(constant_val, torch.nn.Parameter):
+                    log.debug(
+                        "%s created when tracing %s is a parameter. But "
+                        "it's not registered with register_parameter(). export will treat it as a constant tensor",
+                        str(node.target),
+                        str(node.meta.get("stack_trace", "<unknown stack>")),
+                    )
+                    # We get the real data out of the parameter by disabling the surrounding fake mode.
+                    with unset_fake_temporarily():
+                        constant_val = constant_val.data
+                constant_kind = InputKind.CONSTANT_TENSOR
+                constant_fqn = _get_first_fqn(constant_attrs, constant_val)
+                if constant_fqn is not None:
+                    constant_name = constant_fqn.replace(".", "_")
+                else:
+                    constant_name = f"lifted_tensor_{num_tensor_constants}"
+                    constant_fqn = get_constant_fqn(node, constant_name)
+                    while constant_fqn in used_target_names:
+                        num_tensor_constants += 1
+                        constant_name = f"lifted_tensor_{num_tensor_constants}"
+                        constant_fqn = get_constant_fqn(node, constant_name)
+                    num_tensor_constants += 1
+            else:
+                raise SpecViolationError(
+                    f"getattr node {node} referencing unsupported type {type(constant_val)}"
+                )
+
+            with gm.graph.inserting_before(first_user_input):
+                # Insert the constant node before the first user input
+                const_placeholder_node = gm.graph.placeholder(constant_name)
+                # match target name with its node name in case there is name collision
+                # and suffix is added to node name in fx
+                const_placeholder_node.target = const_placeholder_node.name
+
+                for k, v in node.meta.items():
+                    const_placeholder_node.meta[k] = v
+
+                # Once the FQN has been used, remove nn_module_stack, stack_trace
+                const_placeholder_node.meta.pop("nn_module_stack")
+                const_placeholder_node.meta.pop("stack_trace", None)
+
+                input_spec_arg: ArgumentSpec
+                if isinstance(constant_val, torch.Tensor):
+                    if fake_mode is not None:
+                        const_placeholder_node.meta["val"] = fake_mode.from_tensor(
+                            constant_val, static_shapes=True
+                        )
+                        const_placeholder_node.meta["val"].constant = constant_val
+                    else:
+                        const_placeholder_node.meta["val"] = constant_val
+                    input_spec_arg = TensorArgument(name=const_placeholder_node.name)
+                elif isinstance(constant_val, torch._C.ScriptObject):
+                    class_fqn = constant_val._type().qualified_name()  # type: ignore[attr-defined]
+                    const_placeholder_node.meta["val"] = CustomObjArgument(
+                        constant_fqn, class_fqn
+                    )
+                    input_spec_arg = CustomObjArgument(
+                        name=const_placeholder_node.name, class_fqn=class_fqn
+                    )
+                elif isinstance(constant_val, FakeScriptObject):
+                    class_fqn = constant_val.script_class_name
+                    const_placeholder_node.meta["val"] = CustomObjArgument(
+                        constant_fqn, class_fqn, constant_val
+                    )
+                    input_spec_arg = CustomObjArgument(
+                        name=const_placeholder_node.name,
+                        class_fqn=class_fqn,
+                        fake_val=constant_val,
+                    )
+                elif is_opaque_type(type(constant_val)):
+                    class_fqn = get_opaque_type_name(type(constant_val))
+                    fake_val = (
+                        maybe_to_fake_obj(fake_mode, constant_val)
+                        if fake_mode
+                        else None
+                    )
+                    const_placeholder_node.meta["val"] = CustomObjArgument(
+                        constant_fqn,
+                        class_fqn,
+                        fake_val,  # pyrefly: ignore[bad-argument-type]
+                    )
+                    input_spec_arg = CustomObjArgument(
+                        name=const_placeholder_node.name,
+                        class_fqn=class_fqn,
+                        fake_val=fake_val,  # pyrefly: ignore[bad-argument-type]
+                    )
+                else:
+                    raise SpecViolationError(
+                        f"tried to lift unsupported type {type(constant_val)} from node {node.format_node()}"
+                    )
+
+                lifted_objs.add(constant_val, const_placeholder_node)
+                node.replace_all_uses_with(const_placeholder_node)
+                gm.graph.erase_node(node)
+
+                renamed_targets[node.name] = const_placeholder_node.name
+
+                # Add the constant as a buffer to the graph signature
+                graph_signature.input_specs.insert(
+                    first_user_input_loc,
+                    InputSpec(
+                        kind=constant_kind,
+                        arg=input_spec_arg,
+                        target=constant_fqn,
+                    ),
+                )
+                if constant_val in constant_attrs:
+                    for fqn in constant_attrs[constant_val]:
+                        all_constants[fqn] = constant_val
+                else:
+                    all_constants[constant_fqn] = constant_val
+                first_user_input_loc += 1
+
+    for spec in graph_signature.output_specs:
+        if spec.arg.name in renamed_targets:
+            spec.arg.name = renamed_targets[spec.arg.name]
+
+    return all_constants
+
+
+def rewrite_script_object_meta(
+    gm: torch.fx.GraphModule,
+) -> dict[str, _ConstantAttributeType]:
+    """When tracing, we produce a graph with FakeScriptObject in the
+    meta["val"].
+
+    For now, we rewrie meta["val"] to be a placeholder CustomObjArgument
+    """
+    constants: dict[
+        str,
+        _ConstantAttributeType,
+    ] = {}
+    for node in gm.graph.nodes:
+        if "val" not in node.meta:
+            continue
+
+        old_meta = node.meta["val"]
+
+        if isinstance(old_meta, torch.ScriptObject):
+            class_fqn = old_meta._type().qualified_name()  # type: ignore[attr-defined]
+            new_meta = CustomObjArgument(node.name, class_fqn)
+            constants[node.name] = old_meta
+            node.meta["val"] = new_meta
+
+        elif isinstance(old_meta, FakeScriptObject):
+            class_fqn = old_meta.script_class_name  # type: ignore[attr-defined]
+            new_meta = CustomObjArgument(node.name, class_fqn, old_meta)
+            constants[node.name] = old_meta
+            node.meta["val"] = new_meta
+
+    return constants
+
+
+def _materialize_and_lift_constants(
+    gm: torch.fx.GraphModule,
+    export_graph_signature: ExportGraphSignature,
+    constant_attrs: ConstantAttrMap,
+) -> dict[str, _ConstantAttributeType]:
+    constants = rewrite_script_object_meta(gm)
+    constants.update(lift_constants_pass(gm, export_graph_signature, constant_attrs))
+    return constants

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/remove_runtime_assertions.py

@@ -0,0 +1,36 @@
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+
+
+class _RemoveRuntimeAssertionsPass(PassBase):
+    """
+    Remove runtime assertions inserted by the
+    _AddRuntimeAssertionsForInlineConstraintsPass.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        modified = False
+        for module in graph_module.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if node.target in [
+                    torch.ops.aten._assert_async.msg,
+                    torch.ops.aten._assert_scalar.default,
+                    torch.ops.aten.sym_constrain_range_for_size.default,
+                    torch.ops.aten.sym_constrain_range.default,
+                    torch.ops.aten._assert_tensor_metadata.default,
+                ]:
+                    assert_async_node = node
+                    if len(assert_async_node.users) > 0:
+                        continue
+                    module.graph.erase_node(assert_async_node)
+                    # the upstream scalar_tensor <- {le, ge} <- sym_size
+                    # linear chain of nodes of nodes is removed by the
+                    # downstream dead code elimination
+                    modified = True
+
+        # We don't necessarily want to run DCE here because it could affect
+        # nodes that are in the module_call_graph attribute of the exported
+        # program. We will leave it to the pass caller to call DCE.
+        return PassResult(graph_module, modified)

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/replace_autocast_with_hop_pass.py

@@ -0,0 +1,212 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch._higher_order_ops.wrap import wrap_with_autocast
+
+from ..utils import node_inline_, nodes_filter, nodes_first, sequential_split
+from .replace_with_hop_pass_util import (
+    _replace_with_hop_helper,
+    _replace_with_hop_pass_helper,
+    _sequential_split_and_maybe_inline_subgraphs_helper,
+)
+
+
+if TYPE_CHECKING:
+    from torch.export.graph_signature import ExportGraphSignature
+
+
+def _is_autocast_node(node: torch.fx.Node) -> torch.fx.Node | bool:
+    return (
+        node
+        and node.op == "call_function"
+        and node.target
+        in [
+            torch.amp.autocast_mode._enter_autocast,
+            torch.amp.autocast_mode._exit_autocast,
+        ]
+    )
+
+
+def _is_enter_autocast_node(node: torch.fx.Node) -> torch.fx.Node | bool:
+    return (
+        node
+        and node.op == "call_function"
+        and node.target is torch.amp.autocast_mode._enter_autocast
+    )
+
+
+def _is_exit_autocast_node(node: torch.fx.Node) -> torch.fx.Node | bool:
+    return (
+        node
+        and node.op == "call_function"
+        and node.target is torch.amp.autocast_mode._exit_autocast
+    )
+
+
+def _is_autocast_sub_mod(node: torch.fx.Node) -> bool:
+    """
+    Check if the first non-placeholder node is `torch.amp.autocast_mode._enter_autocast`.
+    """
+    if node.op == "call_module":
+        if not isinstance(node.target, str):
+            raise AssertionError(f"expected str target, got {type(node.target)}")
+        subgm = getattr(node.graph.owning_module, node.target)
+        first_non_ph = nodes_first(
+            subgm.graph.nodes, lambda node: node.op != "placeholder"
+        )
+        if (
+            first_non_ph
+            and first_non_ph.op == "call_function"
+            and first_non_ph.target is torch.amp.autocast_mode._enter_autocast
+        ):
+            # TODO: check if current auto-cast type is the same as the args of
+            # _enter_autocast. If so, return False, i.e. do not create a submodule.
+            return True
+    return False
+
+
+def _check_valid_autocast_block(
+    enter_autocast_node: torch.fx.Node, exit_autocast_node: torch.fx.Node
+) -> None:
+    if not _is_enter_autocast_node(enter_autocast_node):
+        raise AssertionError(
+            f"expected enter_autocast node, got {enter_autocast_node.target}"
+        )
+    if not _is_exit_autocast_node(exit_autocast_node):
+        raise AssertionError(
+            f"expected exit_autocast node, got {exit_autocast_node.target}"
+        )
+    if exit_autocast_node.args[0] != enter_autocast_node:
+        raise AssertionError(
+            "exit_autocast_node.args[0] must match enter_autocast_node"
+        )
+
+
+def _replace_with_hop(node: torch.fx.Node) -> None:
+    if node.op != "call_module":
+        raise AssertionError(f"expected call_module op, got {node.op}")
+    graph: torch.fx.Graph = node.graph
+    if graph.owning_module is None:
+        raise AssertionError("graph.owning_module must not be None")
+    gm: torch.fx.GraphModule = graph.owning_module
+    if not isinstance(node.target, str):
+        raise AssertionError(f"expected str target, got {type(node.target)}")
+    sub_gm = getattr(gm, node.target)
+    sub_graph = sub_gm.graph
+    autocast_nodes = nodes_filter(sub_graph.nodes, _is_autocast_node)
+    if len(autocast_nodes) > 0:
+        if len(autocast_nodes) <= 1:
+            raise AssertionError(
+                f"need at least an enter node and an exit node, got {len(autocast_nodes)}"
+            )
+        enter_autocast_node = autocast_nodes[0]
+        exit_autocast_node = autocast_nodes[-1]
+        _check_valid_autocast_block(enter_autocast_node, exit_autocast_node)
+
+        _replace_with_hop_helper(node, enter_autocast_node, wrap_with_autocast)
+        sub_graph.erase_node(exit_autocast_node)
+        sub_graph.erase_node(enter_autocast_node)
+
+
+def _split_autocast(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    """
+    split_autocast creates a new graph module that splits the input graph module into multiple submodules
+    based on the `_enter_autocast` and `_exit_autocast` nodes. It doesn't mutate the input graph module.
+
+    Nodes between the **outer-most** `_enter_autocast` and `_exit_autocast(_enter_autocast)` are split
+    into a submodule. Nested autocast regions are not split.
+    `_enter_autocast` and `_exit_autocast(_enter_autocast)` nodes are in the submodule as well.
+
+    Below is an example of splitting. A, B, C, D, E are blocks of non-autocast nodes in the original graph
+    module. Nodes marked with the same number are grouped into the same submodule.
+    A               # 0
+    enter_autocast  # 1
+    B               # 1
+    exit_autocast   # 1
+    C               # 2
+    enter_autocast  # 3
+    D               # 3
+    exit_autocast   # 3
+    E               # 4
+    """
+    enter_autocast_node_stack: list[torch.fx.Node] = []
+    first_node_after_outer_most_exit: bool = False
+
+    def node_call_back(node: torch.fx.Node) -> bool:
+        nonlocal enter_autocast_node_stack, first_node_after_outer_most_exit
+        increment_id = False
+        if first_node_after_outer_most_exit or (
+            len(enter_autocast_node_stack) == 0 and _is_enter_autocast_node(node)
+        ):
+            if len(enter_autocast_node_stack) != 0:
+                raise AssertionError(
+                    f"expected empty stack, got {len(enter_autocast_node_stack)} items"
+                )
+            first_node_after_outer_most_exit = False
+            increment_id = True
+        if _is_enter_autocast_node(node):
+            enter_autocast_node_stack.append(node)
+        elif _is_exit_autocast_node(node):
+            if len(enter_autocast_node_stack) == 0:
+                raise AssertionError("enter_autocast_node_stack must not be empty")
+            last_enter_autocast_node = enter_autocast_node_stack.pop()
+            if node.args[0] != last_enter_autocast_node:
+                raise AssertionError("exit node args[0] must match last enter node")
+            if len(enter_autocast_node_stack) == 0:
+                # next node should be in the next submodule since
+                # autocast block ends
+                first_node_after_outer_most_exit = True
+        return increment_id
+
+    return sequential_split(gm, node_call_back)
+
+
+def _sequential_split_and_maybe_inline_subgraphs(
+    gm: torch.fx.GraphModule, graph_signature: ExportGraphSignature | None
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Helper function for replace_autocast_with_hop_pass().
+    Split the graph module into multiple subgraphs based on the autocast nodes.
+    For each subgraph, decides whether to construct a HOO subgraph, or inline the calls
+    back into the parent graph module.
+    Nodes between `_enter_autocast` and `_exit_autocast(_enter_autocast)` are considered
+    as a subgraph.
+    """
+    need_replacing = any(_is_autocast_node(node) for node in gm.graph.nodes)
+    if not need_replacing:
+        return gm, graph_signature
+
+    # split_autocast returns a new graph module that could have different output
+    # args names. We need to fix the graph signature in `_sequential_split_and_maybe_inline_subgraphs_helper`.
+    new_gm = _split_autocast(gm)
+
+    def _maybe_inline_or_replace_with_hop(node: torch.fx.Node) -> None:
+        if _is_autocast_sub_mod(node):
+            _replace_with_hop(node)
+        else:
+            if node.op != "call_module":
+                raise AssertionError(f"expected call_module op, got {node.op}")
+            if not isinstance(node.target, str):
+                raise AssertionError(f"expected str target, got {type(node.target)}")
+            node_inline_(node)
+
+    return _sequential_split_and_maybe_inline_subgraphs_helper(
+        new_gm, graph_signature, _maybe_inline_or_replace_with_hop
+    )
+
+
+def replace_autocast_with_hop_pass(
+    gm: torch.fx.GraphModule, graph_signature: ExportGraphSignature | None
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Split gm into sub-graph-modules using `sequential_split_and_maybe_inline_subgraphs`, and
+    then recursively call itself on each of the submodules.
+    """
+    return _replace_with_hop_pass_helper(
+        gm,
+        graph_signature,
+        _sequential_split_and_maybe_inline_subgraphs,
+    )

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py

@@ -0,0 +1,700 @@
+# mypy: allow-untyped-defs
+import logging
+import operator
+
+import torch
+import torch.export._trace
+from torch._ops import OpOverload
+from torch.ao.quantization.fx._decomposed import (
+    dequantize_per_channel,
+    dequantize_per_tensor,
+    quantize_per_tensor,
+)
+from torch.ao.quantization.utils import calculate_qmin_qmax
+from torch.fx.graph_module import _assign_attr
+
+
+log = logging.getLogger(__name__)
+
+# Those values will need to be carried over multiple operators.
+_INPUT_Q_DTYPE: torch.dtype | torch.fx.Node | None = None
+_SCALE: float | torch.fx.Node | None = None
+_ZERO_POINT: float | torch.fx.Node | None = None
+
+
+def int_to_valid_dtype(val: int) -> torch.dtype:
+    from torch._export.converter import _TORCH_ENUM_TO_DTYPE  # No circular import.
+
+    if isinstance(val, torch.dtype):
+        return val
+    dtype = _TORCH_ENUM_TO_DTYPE[val]
+    if dtype == torch.quint8:
+        return torch.uint8
+    elif dtype == torch.qint8:
+        return torch.int8
+    return dtype
+
+
+def fx_enum_to_dtype(gm: torch.fx.GraphModule, val: int) -> torch.fx.Node:
+    return gm.graph.call_function(int_to_valid_dtype, (val,))
+
+
+def insert_quantized_node(
+    gm: torch.fx.GraphModule,
+    val_node: torch.fx.Node,
+    scale_node: float | torch.fx.Node,
+    zero_point_node: float | torch.fx.Node,
+    qmin_node: float | int | torch.fx.Node,
+    qmax_node: float | int | torch.fx.Node,
+    dtype_node: torch.dtype | torch.fx.Node,
+    qscheme: torch.qscheme | None,
+) -> torch.fx.Node:
+    return gm.graph.call_function(
+        quantize_per_tensor,
+        (
+            val_node,
+            scale_node,
+            zero_point_node,
+            qmin_node,
+            qmax_node,
+            dtype_node,
+        ),
+    )
+
+
+def get_dequantized(
+    val: torch.Tensor,
+    scale: float | torch.Tensor,
+    zero_point: float | torch.Tensor,
+    qmin: float | int,
+    qmax: float | int,
+    dtype: torch.dtype,
+    axis: int | None,
+    qscheme: torch.qscheme | None,
+) -> torch.Tensor:
+    if qscheme is torch.per_tensor_affine:
+        return dequantize_per_tensor(
+            val,
+            scale,  # type: ignore[arg-type]
+            zero_point,  # type: ignore[arg-type]
+            qmin,  # type: ignore[arg-type]
+            qmax,  # type: ignore[arg-type]
+            dtype,
+        )
+    elif qscheme is torch.per_channel_affine:
+        return dequantize_per_channel(
+            val,
+            scale,  # type: ignore[arg-type]
+            zero_point,  # type: ignore[arg-type]
+            axis,  # type: ignore[arg-type]
+            qmin,  # type: ignore[arg-type]
+            qmax,  # type: ignore[arg-type]
+            dtype,
+        )
+    else:
+        raise RuntimeError(f"Unsupported dequantization scheme: {qscheme}")
+
+
+def insert_dequantized_node(
+    gm: torch.fx.GraphModule,
+    val_node: torch.fx.Node,
+    scale_node: float | torch.fx.Node,
+    zero_point_node: float | torch.fx.Node,
+    qmin_node: float | int | torch.fx.Node,
+    qmax_node: float | int | torch.fx.Node,
+    dtype_node: torch.dtype | torch.fx.Node,
+    axis_node: int | torch.fx.Node | None,
+    qscheme: torch.qscheme | None,
+) -> torch.fx.Node:
+    if qscheme is torch.per_tensor_affine:
+        return gm.graph.call_function(
+            dequantize_per_tensor,
+            (
+                val_node,
+                scale_node,
+                zero_point_node,
+                qmin_node,
+                qmax_node,
+                dtype_node,
+            ),
+        )
+    elif qscheme is torch.per_channel_affine:
+        return gm.graph.call_function(
+            dequantize_per_channel,
+            (
+                val_node,
+                scale_node,
+                zero_point_node,
+                axis_node,
+                qmin_node,
+                qmax_node,
+                dtype_node,
+            ),
+        )
+    else:
+        raise RuntimeError(f"Unsupported dequantization scheme: {qscheme}")
+
+
+def get_qmin_qmax(dtype: torch.dtype) -> tuple[int | float, int | float]:
+    return calculate_qmin_qmax(None, None, False, dtype, False)  # type: ignore[arg-type]
+
+
+def insert_qmin_qmax_node(
+    gm: torch.fx.GraphModule, dtype_node: torch.dtype | torch.fx.Node
+) -> tuple[torch.fx.Node, torch.fx.Node]:
+    q_min_max_node = gm.graph.call_function(
+        calculate_qmin_qmax, (None, None, False, dtype_node, False)
+    )
+    qmin_node = gm.graph.call_function(operator.getitem, (q_min_max_node, 0))
+    qmax_node = gm.graph.call_function(operator.getitem, (q_min_max_node, 1))
+    return qmin_node, qmax_node
+
+
+def get_script_object(
+    gm: torch.nn.Module, node: torch.fx.Node
+) -> torch._C.ScriptObject:
+    if not isinstance(node, torch.fx.Node):
+        raise AssertionError(f"expected fx.Node, got {type(node).__name__}")
+    if node.op != "get_attr":
+        raise AssertionError(f"expected get_attr op, got {node.op}")
+    attr_name = node.target
+    if not isinstance(attr_name, str):
+        raise AssertionError(f"expected str target, got {type(attr_name).__name__}")
+
+    mod = gm
+    for attr in attr_name.split("."):
+        mod = getattr(mod, attr)
+    if not isinstance(mod, torch._C.ScriptObject):
+        raise AssertionError(f"expected ScriptObject, got {type(mod).__name__}")
+    return mod
+
+
+def insert_weight_and_bias_get_attr_node_from_get_attr_to_scriptobject(
+    gm: torch.fx.GraphModule,
+    param_node: torch.fx.Node,
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
+    """Directly inline tensor from a get_attr fx node."""
+    mod = get_script_object(gm, param_node)
+    w_qtensor, b_qtensor = mod.unpack()  # type: ignore[attr-defined]
+    w_attr_name, b_attr_name = (
+        f"dequantized_{param_node.target}_w",
+        f"dequantized_{param_node.target}_b",
+    )
+    return insert_weight_and_bias_get_attr_node(
+        gm, w_qtensor, b_qtensor, w_attr_name, b_attr_name
+    )
+
+
+def insert_weight_and_bias_get_attr_node_from_get_attr_to_qtensor(
+    gm: torch.fx.GraphModule,
+    get_attr_to_weight_node: torch.fx.Node,
+    get_attr_to_bias_node: torch.fx.Node | None,
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
+    if not isinstance(get_attr_to_weight_node.target, str):
+        raise AssertionError(
+            f"expected str target, got {type(get_attr_to_weight_node.target).__name__}"
+        )
+    w_qtensor = getattr(gm, get_attr_to_weight_node.target)
+    w_attr_name = f"dequantized_{get_attr_to_weight_node.target}_w"
+
+    if get_attr_to_bias_node is not None:
+        if not isinstance(get_attr_to_bias_node.target, str):
+            raise AssertionError(
+                f"expected str target, got {type(get_attr_to_bias_node.target).__name__}"
+            )
+        b_qtensor = getattr(gm, get_attr_to_bias_node.target)
+        b_attr_name = f"dequantized_{get_attr_to_bias_node.target}_b"
+    else:
+        b_qtensor, b_attr_name = None, ""
+
+    return insert_weight_and_bias_get_attr_node(
+        gm, w_qtensor, b_qtensor, w_attr_name, b_attr_name
+    )
+
+
+def insert_weight_and_bias_get_attr_node(
+    gm: torch.fx.GraphModule,
+    w_qtensor: torch.Tensor,
+    b_qtensor: torch.Tensor | None,
+    w_attr_name: str,
+    b_attr_name: str,
+) -> tuple[torch.fx.Node, torch.fx.Node | None]:
+    w_tensor = get_tensor_from_qtensor(w_qtensor)
+    _assign_attr(w_tensor, gm, w_attr_name)
+    w_tensor_attr = gm.graph.get_attr(w_attr_name)
+
+    if b_qtensor is not None:
+        b_tensor = get_tensor_from_qtensor(b_qtensor, dequant=False)
+        _assign_attr(b_tensor, gm, b_attr_name)
+        b_tensor_attr = gm.graph.get_attr(b_attr_name)
+    else:
+        b_tensor_attr = None
+
+    return w_tensor_attr, b_tensor_attr
+
+
+def get_tensor_from_qtensor(
+    qtensor: torch.Tensor, dequant: bool = True
+) -> torch.Tensor:
+    # Manual conversion because qint8 is not used anymore.
+    if qtensor.dtype in [torch.qint8, torch.quint8]:
+        tensor = qtensor.int_repr()
+    else:
+        tensor = qtensor
+
+    # Weights need dequantization with scaling and zero_point adjustment, but
+    # bias does not need that.
+    if dequant:
+        qscheme = qtensor.qscheme()
+        if qscheme == torch.per_channel_affine:
+            scale, zero_point, axis = (
+                qtensor.q_per_channel_scales(),
+                qtensor.q_per_channel_zero_points(),
+                qtensor.q_per_channel_axis(),
+            )
+        else:
+            scale, zero_point, axis = (
+                qtensor.q_scale(),  # type: ignore[assignment]
+                qtensor.q_zero_point(),  # type: ignore[assignment]
+                None,
+            )
+        dtype = tensor.dtype
+        qmin, qmax = get_qmin_qmax(dtype)
+        return get_dequantized(
+            tensor, scale, zero_point, qmin, qmax, dtype, axis, qscheme
+        )
+    return tensor
+
+
+def insert_fused_activation_node(
+    gm: torch.fx.GraphModule, opname: str, fx_node: torch.fx.Node
+) -> torch.fx.Node:
+    if opname in ["conv1d_relu", "conv2d_relu", "linear_relu", "add_relu", "mul_relu"]:
+        fx_node = gm.graph.call_function(torch.ops.aten.relu, (fx_node,))
+    return fx_node
+
+
+def _conv1d_op_with_squeeze(
+    inp: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+) -> torch.Tensor:
+    # In quantized version, conv1d is emulated using conv2d with squeeze and unsqueeze
+    # operations before and after the conv2d operation to match the dimension of weights.
+    # Reference: https://github.com/pytorch/pytorch/blob/eca0cb0fbe84bb0a34fa94afe261bceecd52c436/aten/src/ATen/native/quantized/cpu/qconv.cpp#L1827  # noqa: B950
+    s_inp = torch.ops.aten.unsqueeze(inp, 2)
+    conv1d_res = torch.ops.aten.conv2d(
+        s_inp,
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        groups,
+    )
+    uns_conv1d_res = torch.ops.aten.squeeze(conv1d_res, 2)
+    return uns_conv1d_res
+
+
+def _transform_conv_with_packedparam(gm: torch.fx.GraphModule, node: torch.fx.Node):
+    """Conv specific transformation function."""
+    if not isinstance(node.target, torch._ops.OpOverload):
+        raise AssertionError(f"expected OpOverload, got {type(node.target).__name__}")
+    opname = node.target._opname
+    scale_node, zero_point_node = node.args[2], node.args[3]
+
+    op_f = (
+        torch.ops.aten.conv2d
+        if opname in ["conv2d", "conv2d_relu"]
+        else _conv1d_op_with_squeeze
+    )
+
+    inp_node, param_node = node.args[0], node.args[1]
+    if not isinstance(inp_node, torch.fx.Node):
+        raise AssertionError(f"expected fx.Node for inp, got {type(inp_node)}")
+    if not isinstance(param_node, torch.fx.Node):
+        raise AssertionError(f"expected fx.Node for param, got {type(param_node)}")
+
+    if param_node.op == "call_function":
+        # Using Conv2dPrepackParam from conv_prepack.
+        # We directly skip the packing call and inline weights and bias.
+        w_node, b_node = param_node.args[0], param_node.args[1]
+        if not isinstance(w_node, torch.fx.Node):
+            raise AssertionError(f"expected fx.Node for w, got {type(w_node)}")
+        if b_node is not None and not isinstance(b_node, torch.fx.Node):
+            raise AssertionError(f"expected fx.Node for b, got {type(b_node)}")
+        (
+            param_0,
+            param_1,
+        ) = insert_weight_and_bias_get_attr_node_from_get_attr_to_qtensor(
+            gm, w_node, b_node
+        )
+        op_res_node = gm.graph.call_function(
+            op_f, (inp_node, param_0, param_1, *param_node.args[2:])
+        )
+    else:
+        # Using ConvPrepackedParam.
+        param = get_script_object(gm, param_node)
+        (
+            param_0,
+            param_1,
+        ) = insert_weight_and_bias_get_attr_node_from_get_attr_to_scriptobject(
+            gm, param_node
+        )  # type: ignore[assignment]
+        op_res_node = gm.graph.call_function(
+            op_f,
+            (
+                inp_node,
+                param_0,
+                param_1,
+                param.stride(),  # type: ignore[attr-defined]
+                param.padding(),  # type: ignore[attr-defined]
+                param.dilation(),  # type: ignore[attr-defined]
+                param.groups(),  # type: ignore[attr-defined]
+            ),
+        )
+    return op_res_node, scale_node, zero_point_node
+
+
+def _transform_linear_with_packedparam(gm: torch.fx.GraphModule, node: torch.fx.Node):
+    """Linear specific transformation function."""
+    scale_node, zero_point_node = node.args[2], node.args[3]
+
+    inp_node, param_node = node.args[0], node.args[1]
+    if not isinstance(inp_node, torch.fx.Node):
+        raise AssertionError(f"expected fx.Node for inp, got {type(inp_node)}")
+    if not isinstance(param_node, torch.fx.Node):
+        raise AssertionError(f"expected fx.Node for param, got {type(param_node)}")
+
+    if param_node.op == "call_function":
+        # Using LinearPrepackParam from linear_prepack.
+        # We directly skip the packing call and inline weights and bias.
+        w_node, b_node = param_node.args[0], param_node.args[1]
+        if not isinstance(w_node, torch.fx.Node):
+            raise AssertionError(f"expected fx.Node for w, got {type(w_node)}")
+        if b_node is not None and not isinstance(b_node, torch.fx.Node):
+            raise AssertionError(f"expected fx.Node for b, got {type(b_node)}")
+        (
+            param_0,
+            param_1,
+        ) = insert_weight_and_bias_get_attr_node_from_get_attr_to_qtensor(
+            gm, w_node, b_node
+        )
+        op_res_node = gm.graph.call_function(
+            torch.ops.aten.linear, (inp_node, param_0, param_1, *param_node.args[2:])
+        )
+    else:
+        # Using LinearPackedParams.
+        (
+            param_0,
+            param_1,
+        ) = insert_weight_and_bias_get_attr_node_from_get_attr_to_scriptobject(
+            gm, param_node
+        )  # type: ignore[assignment]
+        op_res_node = gm.graph.call_function(
+            torch.ops.aten.linear, (inp_node, param_0, param_1)
+        )
+    return op_res_node, scale_node, zero_point_node
+
+
+def _transform_op_where_last_two_arguments_are_scale_and_zero_point(
+    gm: torch.fx.GraphModule, node: torch.fx.Node
+):
+    """
+    This transformation function can be used for function where the last two
+    parameters are scale and zero point. Additionally, the function's parameters
+    do not need any unpacking.
+    """
+    to_standard_op = {
+        "mul": torch.ops.aten.mul,
+        "mul_relu": torch.ops.aten.mul,
+        "add": torch.ops.aten.add,
+        "add_relu": torch.ops.aten.add,
+        "softmax": torch.ops.aten.softmax,
+        "cat": torch.ops.aten.cat,
+        "hardswish": torch.ops.aten.hardswish,
+    }
+
+    if not isinstance(node.target, torch._ops.OpOverload):
+        raise AssertionError(f"expected OpOverload, got {type(node.target).__name__}")
+    opname, args = node.target._opname, node.args
+    scale_node, zero_point_node = args[-2], args[-1]
+    op_res_node = gm.graph.call_function(to_standard_op[opname], tuple(args[:-2]))
+    return op_res_node, scale_node, zero_point_node
+
+
+def _transform_scalar_arithmetic(gm: torch.fx.GraphModule, node: torch.fx.Node):
+    """Transform scalar overload for basic arithmetic."""
+    to_standard_op = {
+        "mul": torch.ops.aten.mul.Scalar,
+        "add": torch.ops.aten.add.Scalar,
+    }
+    if not isinstance(node.target, torch._ops.OpOverload):
+        raise AssertionError(f"expected OpOverload, got {type(node.target).__name__}")
+    opname, args = node.target._opname, node.args
+    op_res_node = gm.graph.call_function(to_standard_op[opname], args)
+    return op_res_node, _SCALE, _ZERO_POINT
+
+
+def _transform_prepacked_op(gm: torch.fx.GraphModule, node: torch.fx.Node):
+    """
+    Transformation for functions under prepacked namespace, where they share
+    the same handling logic that [...]OpContext contains all parameters.
+    """
+    if not isinstance(node.target, torch._ops.OpOverload):
+        raise AssertionError(f"expected OpOverload, got {type(node.target).__name__}")
+    opname, args = node.target._opname, node.args
+    op_f = None
+    if opname == "conv2d_clamp_run":
+        op_f = torch.ops.aten.conv2d
+    elif opname == "linear_clamp_run":
+        op_f = torch.ops.aten.linear
+    else:
+        raise RuntimeError(f"Invalid operator {opname}")
+
+    if not isinstance(args[1], torch.fx.Node):
+        raise AssertionError(f"expected fx.Node for args[1], got {type(args[1])}")
+    so = get_script_object(gm, args[1])
+
+    func_args = []
+    func_args += [args[0]]
+    func_args += so.unpack()[:2]  # type: ignore[attr-defined]
+    if opname == "conv2d_clamp_run":
+        func_args += torch.ops.prepacked.unpack_prepacked_sizes_conv2d(so)[2:]
+
+    op_res_node = gm.graph.call_function(op_f, tuple(func_args))
+    return op_res_node
+
+
+def _transform_batch_norm(gm: torch.fx.GraphModule, node: torch.fx.Node):
+    args = node.args
+    scale_node, zero_point_node = args[-2], args[-1]
+    op_res_node = gm.graph.call_function(
+        torch.ops.aten.native_batch_norm, (*args[:-3], False, 0.1, args[-3])
+    )
+    op_res_node = gm.graph.call_function(operator.getitem, (op_res_node, 0))
+    return op_res_node, scale_node, zero_point_node
+
+
+def fx_transform_quantized_op_to_standard_op(
+    gm: torch.fx.GraphModule, node: torch.fx.Node
+) -> torch.fx.Node:
+    global _SCALE, _ZERO_POINT, _INPUT_Q_DTYPE
+
+    if not isinstance(node.target, torch._ops.OpOverload):
+        raise AssertionError(f"expected OpOverload, got {type(node.target).__name__}")
+    opname, overload = node.target._opname, node.target._overloadname
+
+    key = f"{opname}.{overload}"
+    opname_to_transform_f = {
+        "conv1d.new": _transform_conv_with_packedparam,
+        "conv1d_relu.new": _transform_conv_with_packedparam,
+        "conv1d.default": _transform_conv_with_packedparam,
+        "conv1d_relu.default": _transform_conv_with_packedparam,
+        "conv2d.new": _transform_conv_with_packedparam,
+        "conv2d_relu.new": _transform_conv_with_packedparam,
+        "conv2d.default": _transform_conv_with_packedparam,
+        "conv2d_relu.default": _transform_conv_with_packedparam,
+        "linear.default": _transform_linear_with_packedparam,
+        "linear_relu.default": _transform_linear_with_packedparam,
+        "add.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "add_relu.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "mul.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "mul_relu.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "softmax.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "cat.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "hardswish.default": _transform_op_where_last_two_arguments_are_scale_and_zero_point,
+        "batch_norm2d.default": _transform_batch_norm,
+        "mul.Scalar": _transform_scalar_arithmetic,
+        "add.Scalar": _transform_scalar_arithmetic,
+    }
+
+    if f"{key}" not in opname_to_transform_f:
+        raise RuntimeError(f"Unsupported quantized op during transformation: {key}")
+
+    op_res_node, scale_node, zero_point_node = opname_to_transform_f[f"{key}"](gm, node)
+
+    # Add fused activation layer.
+    op_res_node = insert_fused_activation_node(gm, opname, op_res_node)
+    _SCALE, _ZERO_POINT = scale_node, zero_point_node
+
+    if _INPUT_Q_DTYPE is None:
+        raise AssertionError("_INPUT_Q_DTYPE should not be None")
+    qmin_node, qmax_node = insert_qmin_qmax_node(gm, _INPUT_Q_DTYPE)
+    q_fx_node = insert_quantized_node(
+        gm,
+        op_res_node,
+        scale_node,
+        zero_point_node,
+        qmin_node,
+        qmax_node,
+        _INPUT_Q_DTYPE,
+        torch.per_tensor_affine,
+    )
+    dq_fx_node = insert_dequantized_node(
+        gm,
+        q_fx_node,
+        scale_node,
+        zero_point_node,
+        qmin_node,
+        qmax_node,
+        _INPUT_Q_DTYPE,
+        None,
+        torch.per_tensor_affine,
+    )
+    return dq_fx_node
+
+
+def replace_quantized_ops_with_standard_ops(gm: torch.fx.GraphModule):
+    """
+    Replace legacy quantized ops (aten.quantize_per_tensor, quantized.conv) with
+    PT2 ops (quantize_decomposed.quantize_per_tensor, aten.conv).
+
+    Before:    x || -> aten.q        || -> quantized.conv2d     || -> quantized.linear    || -> aten.dq || -> y
+
+    After:     x || -> qd.q -> qd.dq || -> aten.conv2d -> qd.q -> qd.dq || aten.linear -> qd.q -> qd.dq || -> y
+
+    (qd == quantized_decomposed library, q = quantize, dq = dequantize)
+                                          ^
+                                          |
+                getattr(w), getattr(b) from Conv2dParamPrepack
+
+    During each iteration, the transformation spits out the transformed operator, its quantized output,
+    and its dequantized value together. We did this because dequantization need to use the
+    scale and zero point parameters from the quantization to recover the approximate original value. After each
+    iteration, the new dequantization node will be used as the input to the next node (e.g., dq2 -> linear).
+
+    For operators like conv2d and linear, their weights and bias are packed in a quantized format in the ScriptObject.
+    During the transformation, we unpack those objects, get their dequantized tensor, populate those
+    as attributes to the module, and use getattr to access them.
+
+    One exception in the transformation is conv_prepack and linear_prepack. Those calls pack
+    weight and bias constant tensors into ScriptObject, which are then used by subsequent conv2d or linear calls.
+    During transformation, we directly skip transforming conv_prepack or linear_prepack. We check whether ScriptObject to the
+    quantized::conv2d or linear is from conv_prepack or linear_prepack. If it is, we then inline those parameters
+    to the operator by converting them to a getattr fx.node.
+
+    For prepacked::conv2d_clamp_run and prepacked::linear_clamp_run, we directly convert them to aten.conv2d and aten.linear
+    without the need of doing de/quantization.
+
+    Three global variables defined are _INPUT_Q_DTYPE, _SCALE, _ZERO_POINT. _INPUT_Q_DTYPE determines the de/quantization
+    data type, which is the same across the entire program, but it only shows up in the very first quantization
+    call. _SCALE and _ZERO_POINT are used only when operators do not have those specified. E.g., mul.Scalar.
+    """
+
+    global _INPUT_Q_DTYPE
+
+    quantized = False
+
+    last_quantized_node = None
+    # pyrefly: ignore [bad-assignment]
+    for node in gm.graph.nodes:
+        if isinstance(node.target, OpOverload):
+            with gm.graph.inserting_before(node):
+                namespace, opname = node.target.namespace, node.target._opname
+                if namespace == "quantized" and opname not in [
+                    "conv_prepack",
+                    "linear_prepack",
+                ]:
+                    quantized = True
+                    fx_node = fx_transform_quantized_op_to_standard_op(gm, node)
+                    node.replace_all_uses_with(fx_node)
+                    last_quantized_node = fx_node
+                elif namespace == "prepacked":
+                    quantized = True
+                    fx_node = _transform_prepacked_op(gm, node)
+                    node.replace_all_uses_with(fx_node)
+                    last_quantized_node = fx_node
+                elif namespace == "aten" and opname == "quantize_per_tensor":
+                    inp_node, scale_node, zero_point_node, dtype_node = node.args
+                    dtype_node = fx_enum_to_dtype(gm, dtype_node)
+                    _INPUT_Q_DTYPE = dtype_node
+                    qmin_node, qmax_node = insert_qmin_qmax_node(gm, dtype_node)
+                    q_fx_node = insert_quantized_node(
+                        gm,
+                        inp_node,
+                        scale_node,
+                        zero_point_node,
+                        qmin_node,
+                        qmax_node,
+                        dtype_node,
+                        torch.per_tensor_affine,
+                    )
+                    dq_fx_node = insert_dequantized_node(
+                        gm,
+                        q_fx_node,
+                        scale_node,
+                        zero_point_node,
+                        qmin_node,
+                        qmax_node,
+                        dtype_node,
+                        None,
+                        torch.per_tensor_affine,
+                    )
+                    node.replace_all_uses_with(dq_fx_node)
+                    last_quantized_node = dq_fx_node
+                elif namespace == "aten" and opname == "dequantize":
+                    if last_quantized_node is None:
+                        raise AssertionError("last_quantized_node should not be None")
+                    node.replace_all_uses_with(last_quantized_node)
+                else:
+                    last_quantized_node = node
+
+    # Post-processing again to remove legacy ScriptObjects and quantizated tensors
+    # stored as attributes or in the buffer. This is used to clean up the GraphModule
+    # to not trigger tracing errors like missing __obj_flatten__ functions.
+    def _clean_attr(mod: torch.nn.Module):
+        for submod in mod.modules():
+            attr_names_to_clean = set()
+            for k, v in submod.__dict__.items():
+                if isinstance(v, torch.ScriptObject):
+                    attr_names_to_clean.add(k)
+                if k == "_buffers":
+                    buffer_name_to_clean = set()
+
+                    for b_name, b_value in v.items():
+                        if isinstance(b_value, torch.Tensor) and b_value.dtype in [
+                            torch.qint8,
+                            torch.quint8,
+                        ]:
+                            buffer_name_to_clean.add(b_name)
+                    for b_name in buffer_name_to_clean:
+                        v.pop(b_name, None)
+            for attr_name in attr_names_to_clean:
+                delattr(submod, attr_name)
+
+    if quantized:
+        """
+        TODO: SetAttr + quantized ops will result incorrect program. This flag is used to temporarily
+        bypass test cases.
+
+        The deadcode elimination pass is needed to remove legacy quantized ops. Otherwise, retracing
+        will throw errors. However, the current way of SetAttr does inplace update to attributes, so
+        this pass regard them as dead code and remove them. Below is an example of GraphModule before
+        and after the dead code elimination pass.
+
+        class GraphModule(torch.nn.Module):
+            def forward(self, x_1):
+                # No stacktrace found for following nodes
+                data = self.data;  data = None
+                data_1 = self.data
+                add_tensor = torch.ops.aten.add.Tensor(data_1, x_1, alpha = 1);  data_1 = None
+                data_2 = self.data
+                copy_ = torch_Tensor_copy_(data_2, add_tensor);  data_2 = add_tensor = copy_ = None
+                data_3 = self.data
+                add_tensor_1 = torch.ops.aten.add.Tensor(x_1, data_3, alpha = 1);  x_1 = data_3 = None
+                return add_tensor_1
+
+        class GraphModule(torch.nn.Module):
+            def forward(self, x_1):
+                # No stacktrace found for following nodes
+                data_3 = self.data
+                add_tensor_1 = torch.ops.aten.add.Tensor(x_1, data_3, alpha = 1);  x_1 = data_3 = None
+                return add_tensor_1
+        """
+        gm.graph.eliminate_dead_code()
+        _clean_attr(gm)

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/replace_set_grad_with_hop_pass.py

@@ -0,0 +1,131 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch._higher_order_ops.wrap import wrap_with_set_grad_enabled
+
+from ..utils import node_inline_, nodes_filter, nodes_first, nodes_map, sequential_split
+from .replace_with_hop_pass_util import (
+    _replace_with_hop_helper,
+    _replace_with_hop_pass_helper,
+    _sequential_split_and_maybe_inline_subgraphs_helper,
+)
+
+
+if TYPE_CHECKING:
+    from torch.export.graph_signature import ExportGraphSignature
+
+
+def _is_set_grad_enabled_node(node: torch.fx.Node) -> torch.fx.Node | bool:
+    return (
+        node
+        and node.op == "call_function"
+        and node.target is torch._C._set_grad_enabled
+    )
+
+
+def _is_set_grad_enabled_sub_mod(
+    node: torch.fx.Node, omit_if_same_with_ambient: bool = False
+) -> bool | torch.Tensor:
+    if node.op == "call_module":
+        if not isinstance(node.target, str):
+            raise AssertionError(f"expected str target, got {type(node.target)}")
+        subgm = getattr(node.graph.owning_module, node.target)
+        first_non_ph = nodes_first(
+            subgm.graph.nodes, lambda node: node.op != "placeholder"
+        )
+        if (
+            first_non_ph
+            and first_non_ph.op == "call_function"
+            and first_non_ph.target is torch._C._set_grad_enabled
+        ):
+            return (
+                first_non_ph.args[0] != torch.is_grad_enabled()
+                if omit_if_same_with_ambient
+                else True
+            )
+    return False
+
+
+def _replace_with_hop(node: torch.fx.Node) -> None:
+    if node.op != "call_module":
+        raise AssertionError(f"expected call_module op, got {node.op}")
+    graph: torch.fx.Graph = node.graph
+    if graph.owning_module is None:
+        raise AssertionError("graph.owning_module must not be None")
+    gm: torch.fx.GraphModule = graph.owning_module
+    if not isinstance(node.target, str):
+        raise AssertionError(f"expected str target, got {type(node.target)}")
+    sub_gm = getattr(gm, node.target)
+    sub_graph = sub_gm.graph
+    set_grad_nodes = nodes_filter(sub_graph.nodes, _is_set_grad_enabled_node)
+    if len(set_grad_nodes) > 0:
+        if len(set_grad_nodes) != 1:
+            raise AssertionError(
+                f"expected exactly 1 set_grad node, got {len(set_grad_nodes)}"
+            )
+        set_grad_node = set_grad_nodes[0]
+        _replace_with_hop_helper(node, set_grad_node, wrap_with_set_grad_enabled)
+        sub_graph.erase_node(set_grad_node)
+
+
+def _remove_set_grad_and_inline(node: torch.fx.Node) -> None:
+    if node.op != "call_module":
+        raise AssertionError(f"expected call_module op, got {node.op}")
+    graph: torch.fx.Graph = node.graph
+    if graph.owning_module is None:
+        raise AssertionError("graph.owning_module must not be None")
+    gm: torch.fx.GraphModule = graph.owning_module
+    if not isinstance(node.target, str):
+        raise AssertionError(f"expected str target, got {type(node.target)}")
+    sub_gm = getattr(gm, node.target)
+    sub_graph = sub_gm.graph
+    nodes_map(
+        sub_graph.nodes,
+        lambda n: sub_graph.erase_node(n) if _is_set_grad_enabled_node(n) else n,
+    )
+    node_inline_(node)
+
+
+def _sequential_split_and_maybe_inline_subgraphs(
+    gm: torch.fx.GraphModule, graph_signature: ExportGraphSignature | None
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Helper function for replace_set_grad_with_hop_pass().
+    Split the graph module into multiple subgraphs based on the set_grad_enabled nodes.
+    For each subgraph, decides whether to construct a HOO subgraph, or inline the calls
+    back into the parent graph module.
+    """
+    need_replacing = any(_is_set_grad_enabled_node(node) for node in gm.graph.nodes)
+    if not need_replacing:
+        return gm, graph_signature
+
+    # sequential_split returns a new graph module that could have different output
+    # args names. We need to fix the graph signature.
+    new_gm = sequential_split(gm, _is_set_grad_enabled_node)
+
+    def _maybe_inline_or_replace_with_hop(node: torch.fx.Node):
+        if _is_set_grad_enabled_sub_mod(node, omit_if_same_with_ambient=True):
+            _replace_with_hop(node)
+        else:
+            _remove_set_grad_and_inline(node)
+
+    return _sequential_split_and_maybe_inline_subgraphs_helper(
+        new_gm, graph_signature, _maybe_inline_or_replace_with_hop
+    )
+
+
+def replace_set_grad_with_hop_pass(
+    gm: torch.fx.GraphModule, graph_signature: ExportGraphSignature | None
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Split gm into sub-graph-modules using `sequential_split_and_maybe_inline_subgraphs`, and
+    then recursively call itself on each of the submodules.
+    """
+    return _replace_with_hop_pass_helper(
+        gm,
+        graph_signature,
+        _sequential_split_and_maybe_inline_subgraphs,
+    )

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py

@@ -0,0 +1,64 @@
+# mypy: allow-untyped-defs
+
+import torch
+from torch._export.error import InternalError
+from torch._export.pass_base import _ExportPassBaseDeprecatedDoNotUse
+from torch._ops import HigherOrderOperator, OpOverload
+
+
+__all__ = ["ReplaceViewOpsWithViewCopyOpsPass"]
+
+
+_NON_FUNCTIONAL_OPS_TO_FUNCTIONAL_OPS: dict[OpOverload, OpOverload] = {
+    torch.ops.aten._unsafe_view.default: torch.ops.aten.view_copy.default,
+}
+
+
+def is_view_op(schema: torch._C.FunctionSchema) -> bool:
+    if len(schema.arguments) == 0:
+        return False
+    alias_info = schema.arguments[0].alias_info
+    return (alias_info is not None) and (not alias_info.is_write)
+
+
+def get_view_copy_of_view_op(schema: torch._C.FunctionSchema) -> OpOverload | None:
+    if is_view_op(schema) and schema.name.startswith("aten::"):
+        view_op_name = schema.name.split("::")[1]
+        view_op_overload = (
+            schema.overload_name if schema.overload_name != "" else "default"
+        )
+        view_copy_op_name = view_op_name + "_copy"
+        if not hasattr(torch.ops.aten, view_copy_op_name):
+            raise InternalError(f"{schema.name} is missing a view_copy variant")
+
+        view_copy_op_overload_packet = getattr(torch.ops.aten, view_copy_op_name)
+
+        if not hasattr(view_copy_op_overload_packet, view_op_overload):
+            raise InternalError(f"{schema.name} is missing a view_copy variant")
+
+        return getattr(view_copy_op_overload_packet, view_op_overload)
+
+    return None
+
+
+class ReplaceViewOpsWithViewCopyOpsPass(_ExportPassBaseDeprecatedDoNotUse):
+    """
+    Our backend expects pure functional operators. For efficiency
+    purposes, we keep view ops around while functionalizing the exported
+    program. This pass replaces view ops with view copy ops for backends that
+    need AOT memory planning.
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op in _NON_FUNCTIONAL_OPS_TO_FUNCTIONAL_OPS:
+            return super().call_operator(
+                (_NON_FUNCTIONAL_OPS_TO_FUNCTIONAL_OPS[op]), args, kwargs, meta
+            )
+
+        if isinstance(op, HigherOrderOperator):
+            return super().call_operator(op, args, kwargs, meta)
+
+        if view_copy_op := get_view_copy_of_view_op(op._schema):
+            return super().call_operator(view_copy_op, args, kwargs, meta)
+
+        return super().call_operator(op, args, kwargs, meta)

kronos-venv/lib/python3.12/site-packages/torch/_export/passes/replace_with_hop_pass_util.py

@@ -0,0 +1,202 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import contextlib
+import copy
+import operator
+from typing import TYPE_CHECKING
+
+import torch
+
+from ..utils import node_replace_, nodes_map
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch._ops import HigherOrderOperator
+    from torch.export.graph_signature import ExportGraphSignature
+
+
+def _replace_with_hop_helper(
+    node: torch.fx.Node,
+    enter_block_node: torch.fx.Node,
+    wrap_hoo: HigherOrderOperator,
+) -> None:
+    graph: torch.fx.Graph = node.graph
+    if graph.owning_module is None:
+        raise AssertionError("graph.owning_module must not be None")
+    gm: torch.fx.GraphModule = graph.owning_module
+    if not isinstance(node.target, str):
+        raise AssertionError(f"expected str target, got {type(node.target)}")
+    sub_gm = getattr(gm, node.target)
+
+    def set_hoo_node_meta(call_func_node):
+        call_func_node.meta["nn_module_stack"] = copy.copy(
+            enter_block_node.meta.get("nn_module_stack", {})
+        )
+        call_func_node.meta["torch_fn"] = (
+            f"{wrap_hoo.__name__}",
+            # pyrefly: ignore [missing-attribute]
+            f"{wrap_hoo.__class__.__name__}.{wrap_hoo.__name__}",
+        )
+        if isinstance(output_args, (tuple, list)):
+            call_func_node.meta["val"] = tuple(arg.meta["val"] for arg in output_args)
+        elif isinstance(output_args, torch.fx.Node):
+            call_func_node.meta["val"] = (output_args.meta["val"],)
+
+    with graph.inserting_before(node):
+        get_attr_node = graph.get_attr(node.target)
+        get_attr_node.meta["nn_module_stack"] = copy.copy(
+            enter_block_node.meta.get("nn_module_stack", {})
+        )
+        output_node = next(iter(reversed(sub_gm.graph.nodes)), None)
+        # Split_module pass intentionally doesn't add output node
+        # if the graph doesn't return anything.
+        # TODO (tmanlaibaatar) Figure out if this is right behaviour
+        # for split_module
+        if isinstance(output_node, torch.fx.Node) and output_node.op != "output":
+            output_node = None
+        if output_node is not None:
+            if len(output_node.args) != 1:
+                raise AssertionError(
+                    f"expected 1 output arg, got {len(output_node.args)}"
+                )
+            output_args = output_node.args[0]
+            enter_block_node_args = enter_block_node.args
+            if isinstance(output_args, (tuple, list)):
+                call_func_node = graph.call_function(
+                    wrap_hoo,
+                    (*enter_block_node_args, get_attr_node, *node.args),
+                    {},
+                )
+                # Create the metadata
+                set_hoo_node_meta(call_func_node)
+                node_replace_(node, call_func_node)
+
+                # Rename the name of getitem nodes to the actual name of its contents
+                # for passing verifier and better readability, also propagate metadata
+                for get_item_node in call_func_node.users:
+                    idx: int = get_item_node.args[1]  # type: ignore[assignment]
+                    output_node = output_args[idx]
+                    get_item_node._rename(output_node.name)
+                    get_item_node.meta = output_node.meta
+
+            elif isinstance(output_args, torch.fx.Node):
+                call_func_node = graph.create_node(
+                    "call_function",
+                    wrap_hoo,
+                    (*enter_block_node_args, get_attr_node, *node.args),
+                    {},
+                    output_args.name,
+                )
+                # Modify the subgraph to output a singleton list.
+                output_node.args = ((output_args,),)
+                # Add in an extra `getitem(wrap_hoo, 0)` node to the toplevel graph.
+                get_item_node = graph.create_node(
+                    "call_function",
+                    operator.getitem,
+                    (call_func_node, 0),
+                    {},
+                )
+                # Create the metadata
+                get_item_node.meta = output_args.meta
+                set_hoo_node_meta(call_func_node)
+                node_replace_(node, get_item_node)
+            else:
+                raise NotImplementedError(
+                    f"replace_with_hop_pass doesn't support output type {type(output_args)}"
+                )
+        else:
+            # TODO (shangdiy): remove this line, since the export graph can be non-functional
+            node.graph.erase_node(node)
+
+
+def _sequential_split_and_maybe_inline_subgraphs_helper(
+    new_gm: torch.fx.GraphModule,
+    graph_signature: ExportGraphSignature | None,
+    maybe_inline_or_replace_with_hop: Callable[[torch.fx.Node], None],
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Helper function for replacing graph nodse with higher order nodes.
+    For each subgraph in `new_gm`, decides whether to construct a HOO subgraph, or inline the calls
+    back into the parent graph module, depending on `maybe_inline_or_replace_with_hop`.
+    """
+    # new_gm is a new graph module that could have different output args names.
+    # We need to fix the graph signature.
+    replace_ctx = contextlib.nullcontext()
+    new_signature = None
+    if graph_signature is not None:
+        # Cannot deep copy a real ScriptObject, which is referenced
+        # in the FakeScriptObject. Copy should be good enough to guard
+        # against accidental mutation to original graph_signature.
+        new_signature = copy.copy(graph_signature)
+        new_gm_out_node = next(reversed(new_gm.graph.find_nodes(op="output")))
+        if new_gm_out_node.op != "output" or len(new_gm_out_node.args[0]) != len(
+            new_signature.output_specs
+        ):
+            raise AssertionError(
+                f"output node mismatch: {new_gm_out_node.op}, "
+                f"{len(new_gm_out_node.args[0])} vs {len(new_signature.output_specs)}"
+            )
+        for arg_node, out_spec in zip(
+            new_gm_out_node.args[0], new_signature.output_specs
+        ):
+            if arg_node is None:
+                if out_spec.arg.value is not None:  # type: ignore[union-attr]
+                    raise AssertionError(
+                        f"expected None out_spec.arg.value, got {out_spec.arg.value}"  # type: ignore[union-attr]
+                    )
+            elif (
+                isinstance(arg_node, torch.fx.Node)
+                and out_spec.arg.name != arg_node.name
+            ):
+                out_spec.arg.name = arg_node.name
+
+        replace_ctx = new_gm._set_replace_hook(new_signature.get_replace_hook())  # type: ignore[assignment]
+
+    with replace_ctx:
+        nodes_map(
+            list(new_gm.graph.nodes),
+            lambda node: (
+                maybe_inline_or_replace_with_hop(node)
+                if node.op == "call_module"
+                else node
+            ),
+        )
+    new_gm.recompile()
+    new_gm.graph.lint()
+    return new_gm, new_signature
+
+
+def _replace_with_hop_pass_helper(
+    gm: torch.fx.GraphModule,
+    graph_signature: ExportGraphSignature | None,
+    sequential_split_and_maybe_inline_subgraphs: Callable[
+        [torch.fx.GraphModule, ExportGraphSignature | None],
+        tuple[torch.fx.GraphModule, ExportGraphSignature | None],
+    ],
+) -> tuple[torch.fx.GraphModule, ExportGraphSignature | None]:
+    """
+    Split gm into sub-graph-modules using `sequential_split_and_maybe_inline_subgraphs`, and
+    then recursively call itself on each of the submodules.
+    """
+    new_gm, new_signature = sequential_split_and_maybe_inline_subgraphs(
+        gm, graph_signature
+    )
+    # recursively call
+    for node in new_gm.graph.nodes:
+        if node.op == "get_attr":
+            subgm = getattr(new_gm, node.target)
+            if not isinstance(subgm, torch.fx.GraphModule):
+                continue
+            new_subgm, _ = _replace_with_hop_pass_helper(
+                subgm,
+                None,
+                sequential_split_and_maybe_inline_subgraphs,
+            )
+            setattr(new_gm, node.target, new_subgm)
+
+    new_gm.recompile()
+    new_gm.graph.lint()
+    return new_gm, new_signature

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/dynamic_shapes.py

@@ -0,0 +1,327 @@
+import dataclasses
+from typing import Any
+
+import torch
+from torch._dynamo.exc import UserError, UserErrorType
+from torch.export.dynamic_shapes import (
+    _check_dynamic_shapes,
+    _DerivedDim,
+    _DimHint,
+    _tree_map_with_path,
+    Dim,
+)
+from torch.utils._pytree import tree_map
+
+from .serialize import _dataclass_to_dict
+
+
+@dataclasses.dataclass
+class RootDim:
+    """
+    This represents a Dim object.
+    """
+
+    min: int
+    max: int | None
+    derived: list[str]
+
+
+@dataclasses.dataclass
+class DynamicShapesSpec:
+    """
+    This stores a dynamic_shapes spec for de/serialization.
+    """
+
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None
+    dims: dict[str, RootDim]
+
+
+def _postprocess_serialized_shapes(
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
+    dims: dict[str, dict[str, int | list[str] | None]],
+    to_dict: bool | None = False,
+) -> DynamicShapesSpec | dict[str, Any]:
+    """
+    Sorts dims and dumps to dictionary format.
+    """
+    from torch.utils._sympy.numbers import int_oo
+
+    # pyrefly: ignore [bad-assignment]
+    dims = {
+        k: RootDim(
+            min=v["min"],  # type: ignore[arg-type]
+            max=None if v["max"] is int_oo else v["max"],  # type: ignore[arg-type]
+            derived=sorted(v["derived"]),  # type: ignore[arg-type]
+        )
+        for k, v in sorted(dims.items())
+    }
+    # pyrefly: ignore [bad-argument-type]
+    spec = DynamicShapesSpec(dynamic_shapes=dynamic_shapes, dims=dims)
+    if to_dict:
+        return _dataclass_to_dict(spec)
+    else:
+        return spec
+
+
+def _dump_dynamic_shapes(
+    dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
+    args: tuple[Any],
+    kwargs: dict[str, Any] | None = None,
+    to_dict: bool | None = False,
+) -> DynamicShapesSpec | dict[str, Any]:
+    """
+    Utility function for dynamic shapes serialization, serializing a dynamic_shapes spec.
+    Returns a DynamicShapesSpec dataclass containing 2 fields, "dynamic_shapes" and "dims".
+    Uses args & kwargs to distinguish between tensor-level and dim-level specs (only for Nones).
+
+    dynamic_shapes: A pytree structure mirroring the dynamic_shapes input to export():
+        - Each tensor input is represented with a list of values, non-tensor inputs with None.
+        - dynamic dimensions (i.e. symbols) in tensors and Dim enums are represented with strings.
+        - static dimensions are represented with ints.
+
+    dims: A dictionary mapping each symbol name to the min/max range and derived dim names.
+
+    For example:
+    ```
+    dx = Dim("dx", min=4, max=16)
+    dy = dx + 1
+
+    inputs = (
+        [
+            torch.randn(4, 4),
+            torch.randn(5, 4),
+        ],
+        torch.randn(4),
+        torch.randn(4, 4),
+        "hello",
+    )
+    dynamic_shapes = {
+        "a": [
+            (dx, 4),
+            (dy, 4),
+        ],
+        "b": (Dim.STATIC,),
+        "c": None,
+        "d": None,
+    }
+    out = _dump_dynamic_shapes(dynamic_shapes, inputs, to_dict=True)
+    ```
+    would generate the following output:
+    ```
+    {
+        "dynamic_shapes": (
+            [
+                ["dx", 4],
+                ["dx + 1", 4],
+            ],
+            ["_DimHint.STATIC"],
+            ["_DimHint.STATIC", "_DimHint.STATIC"],
+            None,
+        ),
+        "dims": {
+            "dx": {
+                "min": 4,
+                "max": 16,
+                "derived": ["dx + 1"],
+            },
+        },
+    }
+    ```
+    """
+    dims: dict[str, dict[str, Any]] = {}
+
+    def _standardize_shapes(path, tensor, shape):  # type: ignore[no-untyped-def]
+        """
+        Helps standardize the dynamic_shapes tree structure we serialize,
+        returning lists for each tensor shape, handling tensor-level Nones.
+        """
+        if not isinstance(tensor, torch.Tensor):
+            return None
+        if shape is None:
+            return [Dim.STATIC] * len(tensor.shape)
+
+        out = []
+        if isinstance(shape, dict):
+            for i, s in enumerate(tensor.shape):
+                out.append(s if shape.get(i) is None else shape.get(i))
+        else:
+            if not isinstance(shape, (tuple, list)):
+                raise AssertionError(f"expected tuple or list, got {type(shape)}")
+            for i, s in enumerate(tensor.shape):
+                out.append(s if shape[i] is None else shape[i])
+        return out
+
+    def _track_dim_from_dims(
+        val: None | int | _DimHint | Dim,
+    ) -> None | int | str:
+        """
+        Tracks dims, ranges, derived dims from the standardized dynamic_shapes spec.
+        """
+        if val is None or isinstance(val, int):  # non-tensor input or static
+            return val
+        if isinstance(val, _DimHint):  # store enum as string
+            return val.__class__.__name__ + "." + val.type.name
+
+        if not isinstance(val, Dim):
+            raise AssertionError(f"expected Dim, got {type(val)}")
+
+        # track root dim
+        root = val.root if isinstance(val, _DerivedDim) else val  # type: ignore[attr-defined]
+        if root.__name__ not in dims:
+            dims[root.__name__] = {
+                "min": root.min,  # type: ignore[attr-defined,union-attr]
+                "max": root.max,  # type: ignore[attr-defined,union-attr]
+                "derived": set(),
+            }
+
+        # track derived dims
+        if isinstance(val, _DerivedDim):
+            dims[root.__name__]["derived"].add(val.__name__)
+
+        return val.__name__
+
+    if dynamic_shapes is None:
+        return {"dynamic_shapes": None, "dims": {}}
+
+    # convert to tuple of specs, for each arg/kwarg
+    kwargs = kwargs or {}
+    if isinstance(dynamic_shapes, dict):
+        dynamic_shapes = dynamic_shapes.values()  # type: ignore[assignment]
+    # pyrefly: ignore [bad-assignment, bad-argument-type]
+    dynamic_shapes = tuple(dynamic_shapes)
+    combined_args = tuple(args) + tuple(kwargs.values())
+
+    # run same check when we're processing shapes for export - is this too lazy?
+    _check_dynamic_shapes(dict(enumerate(combined_args)), dynamic_shapes)  # type: ignore[arg-type]
+
+    tree_shapes = _tree_map_with_path(
+        _standardize_shapes, combined_args, dynamic_shapes, tree_name="inputs"
+    )
+    serialized_shapes = tree_map(_track_dim_from_dims, tree_shapes)
+    return _postprocess_serialized_shapes(serialized_shapes, dims, to_dict=to_dict)
+
+
+def _load_dynamic_shapes(
+    spec: DynamicShapesSpec | dict[str, Any],
+    from_dict: bool | None = False,
+) -> dict[str, Any] | tuple[Any] | list[Any] | None:
+    """
+    Utility function for dynamic shapes serialization.
+    Deserializes a DynamicShapesSpec or corresponding dictionary into a dynamic_shapes input to export().
+    """
+    import sympy
+
+    from torch.fx.experimental.symbolic_shapes import _is_supported_equivalence
+
+    if from_dict:
+        if not isinstance(spec, dict):
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                f"With from_dict=True, expected `spec` to be a dict, got {type(spec)}",
+            )
+        if sorted(spec.keys()) != ["dims", "dynamic_shapes"]:
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                "With from_dict=True, expected `spec` to have keys `dims` and `dynamic_shapes`, "
+                f"instead found {spec.keys()}",
+            )
+        dims = {}
+        for k, v in spec["dims"].items():
+            if not isinstance(k, str):
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected `spec['dims']` keys to be strings for symbols, got key {type(k)}",
+                )
+            if sorted(v.keys()) != ["derived", "max", "min"]:
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected `spec['dims']` values to have keys `derived`, `max`, and `min`, "
+                    f"instead found {v.keys()}",
+                )
+            if not isinstance(v["min"], int):
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected dims in `spec['dims']` to map `min` to an int, got {k}: {v['min']}",
+                )
+            if not isinstance(v["max"], int) or v["max"] is None:
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected dims in `spec['dims']` to map `max` to an int or None, got {k}: {v['max']}",
+                )
+            if not isinstance(v["derived"], list) or any(
+                not isinstance(d, str) for d in v["derived"]
+            ):
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    "Expected dims in `spec['dims']` to map `derived` to a list of derived expressions, "
+                    f"got {k}: {v['derived']}",
+                )
+            dims[k] = RootDim(**v)
+        dynamic_shapes = spec["dynamic_shapes"]
+    else:
+        if not isinstance(spec, DynamicShapesSpec):
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                f"Expected `spec` to be a DynamicShapesSpec, got {type(spec)}",
+            )
+        dims = spec.dims
+        dynamic_shapes = spec.dynamic_shapes
+
+    if dynamic_shapes is None:
+        return None
+
+    dim_cache = {}
+    for name, info in dims.items():
+        symbol = sympy.sympify(name)
+        if not isinstance(symbol, sympy.Symbol):
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                f"Expected `spec['dims']` keys to be symbols, got {name}",
+            )
+        dim_cache[name] = Dim(name, min=info.min, max=info.max)  # cache root dim
+        for _expr in info.derived:
+            expr = sympy.sympify(_expr)
+            if len(expr.free_symbols) != 1 or symbol not in expr.free_symbols:
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected derived expressions in to have {name} as the only free symbol, got {expr}",
+                )
+            if not _is_supported_equivalence(expr):
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"Expected derived expressions to be linear expressions, got {expr}",
+                )
+            modulus, remainder = sympy.polys.polytools.div(expr, symbol)
+            ddim = dim_cache[name]
+            if modulus != 1:
+                ddim = int(modulus) * ddim  # type: ignore[assignment, operator]
+            if remainder != 0:
+                ddim = ddim + int(remainder)  # type: ignore[assignment, operator]
+            dim_cache[_expr] = ddim  # cache derived dims
+
+    def deserialize_shape(
+        val: None | int | str,
+    ) -> None | int | Dim | _DimHint:
+        if val is None or isinstance(val, int):
+            return val
+        elif val == "_DimHint.AUTO":
+            return _DimHint.AUTO()
+        elif val == "_DimHint.DYNAMIC":
+            return _DimHint.DYNAMIC()
+        elif val == "_DimHint.STATIC":
+            return _DimHint.STATIC()
+        if not isinstance(val, str):
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                "Expected leaves in `spec['dynamic_shapes']` to be ints, None, Dim.AUTO/STATIC, symbols, "
+                f" or derived expressions, got {val}",
+            )
+        if val not in dim_cache:
+            raise UserError(
+                UserErrorType.INVALID_INPUT,
+                "Expected dims in `spec['dynamic_shapes']` to be tracked in `spec['dims']`, "
+                f"got {val} which is not in {dims.keys()}",
+            )
+        return dim_cache[val]  # type: ignore[return-value]
+
+    return tree_map(deserialize_shape, dynamic_shapes)

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/export_schema.thrift

@@ -0,0 +1,383 @@
+// @generated by update_schema.py
+// checksum<<ffbfd7b406d10b13faddd5faafc7296f90cf65af9409b166c135f6169946a08c>>
+
+namespace py3 torch._export
+namespace cpp2 torch._export.schema
+
+enum ArgumentKind {
+  UNKNOWN = 0,
+  POSITIONAL = 1,
+  KEYWORD = 2,
+}
+
+
+enum Layout {
+  Unknown = 0,
+  SparseCoo = 1,
+  SparseCsr = 2,
+  SparseCsc = 3,
+  SparseBsr = 4,
+  SparseBsc = 5,
+  _mkldnn = 6,
+  Strided = 7,
+}
+
+
+enum MemoryFormat {
+  Unknown = 0,
+  ContiguousFormat = 1,
+  ChannelsLast = 2,
+  ChannelsLast3d = 3,
+  PreserveFormat = 4,
+}
+
+
+enum ScalarType {
+  UNKNOWN = 0,
+  BYTE = 1,
+  CHAR = 2,
+  SHORT = 3,
+  INT = 4,
+  LONG = 5,
+  HALF = 6,
+  FLOAT = 7,
+  DOUBLE = 8,
+  COMPLEXHALF = 9,
+  COMPLEXFLOAT = 10,
+  COMPLEXDOUBLE = 11,
+  BOOL = 12,
+  BFLOAT16 = 13,
+  UINT16 = 28,
+  FLOAT8E4M3FN = 29,
+  FLOAT8E5M2 = 30,
+  FLOAT8E4M3FNUZ = 31,
+  FLOAT8E5M2FNUZ = 32,
+  FLOAT8E8M0FNU = 33,
+  UINT32 = 34,
+  UINT64 = 35,
+}
+
+
+struct Device {
+  10: string type;
+  20: optional i64 index;
+}
+
+union SymExprHint {
+  10: i64 as_int;
+  20: bool as_bool;
+  30: double as_float;
+}
+
+struct SymExpr {
+  10: string expr_str;
+  20: optional SymExprHint hint;
+}
+
+union SymInt {
+  10: SymExpr as_expr;
+  20: i64 as_int;
+}
+
+union SymFloat {
+  10: SymExpr as_expr;
+  20: double as_float;
+}
+
+union SymBool {
+  10: SymExpr as_expr;
+  20: bool as_bool;
+}
+
+struct TensorMeta {
+  10: ScalarType dtype;
+  20: list<SymInt> sizes;
+  30: bool requires_grad;
+  40: Device device;
+  50: list<SymInt> strides;
+  60: SymInt storage_offset;
+  70: Layout layout;
+}
+
+union SymIntArgument {
+  10: string as_name;
+  20: i64 as_int;
+}
+
+union SymFloatArgument {
+  10: string as_name;
+  20: double as_float;
+}
+
+union SymBoolArgument {
+  10: string as_name;
+  20: bool as_bool;
+}
+
+struct TensorArgument {
+  10: string name;
+}
+
+struct TokenArgument {
+  10: string name;
+}
+
+union OptionalTensorArgument {
+  20: TensorArgument as_tensor;
+  10: bool as_none;
+}
+
+struct GraphArgument {
+  10: string name;
+  20: Graph graph;
+}
+
+struct CustomObjArgument {
+  10: string name;
+  20: string class_fqn;
+}
+
+struct ComplexValue {
+  10: double real;
+  20: double imag;
+}
+
+union Argument {
+  10: bool as_none;
+  20: TensorArgument as_tensor;
+  30: list<TensorArgument> as_tensors;
+  50: i64 as_int;
+  70: list<i64> as_ints;
+  80: double as_float;
+  90: list<double> as_floats;
+  100: string as_string;
+  101: list<string> as_strings;
+  110: SymIntArgument as_sym_int;
+  120: list<SymIntArgument> as_sym_ints;
+  130: ScalarType as_scalar_type;
+  140: MemoryFormat as_memory_format;
+  150: Layout as_layout;
+  160: Device as_device;
+  170: bool as_bool;
+  180: list<bool> as_bools;
+  182: SymBoolArgument as_sym_bool;
+  184: list<SymBoolArgument> as_sym_bools;
+  200: GraphArgument as_graph;
+  190: list<OptionalTensorArgument> as_optional_tensors;
+  210: CustomObjArgument as_custom_obj;
+  220: string as_operator;
+  230: SymFloatArgument as_sym_float;
+  240: list<SymFloatArgument> as_sym_floats;
+  250: OptionalTensorArgument as_optional_tensor;
+  260: ComplexValue as_complex;
+  270: list<list<TensorArgument>> as_nested_tensors;
+  280: list<list<i64>> as_int_lists;
+  290: map<string, Argument> as_string_to_argument;
+  300: list<list<double>> as_float_lists;
+}
+
+struct NamedArgument {
+  10: string name;
+  20: Argument arg;
+  30: optional ArgumentKind kind;
+}
+
+struct Node {
+  10: string target;
+  20: list<NamedArgument> inputs;
+  30: list<Argument> outputs;
+  40: map<string, string> metadata;
+  50: optional bool is_hop_single_tensor_return;
+  60: optional string name;
+}
+
+struct Graph {
+  10: list<Argument> inputs;
+  20: list<Argument> outputs;
+  30: list<Node> nodes;
+  40: map<string, TensorMeta> tensor_values;
+  50: map<string, SymInt> sym_int_values;
+  60: map<string, SymBool> sym_bool_values;
+  70: bool is_single_tensor_return;
+  80: map<string, CustomObjArgument> custom_obj_values;
+  90: map<string, SymFloat> sym_float_values;
+}
+
+struct UserInputSpec {
+  10: Argument arg;
+}
+
+union ConstantValue {
+  10: bool as_none;
+  20: i64 as_int;
+  30: double as_float;
+  40: string as_string;
+  50: bool as_bool;
+}
+
+struct InputToConstantInputSpec {
+  10: string name;
+  20: ConstantValue value;
+}
+
+struct InputToParameterSpec {
+  10: TensorArgument arg;
+  20: string parameter_name;
+}
+
+struct InputToBufferSpec {
+  10: TensorArgument arg;
+  20: string buffer_name;
+  30: bool persistent;
+}
+
+struct InputToTensorConstantSpec {
+  10: TensorArgument arg;
+  20: string tensor_constant_name;
+}
+
+struct InputToCustomObjSpec {
+  10: CustomObjArgument arg;
+  20: string custom_obj_name;
+}
+
+struct InputTokenSpec {
+  10: TokenArgument arg;
+}
+
+union InputSpec {
+  10: UserInputSpec user_input;
+  20: InputToParameterSpec parameter;
+  30: InputToBufferSpec buffer;
+  40: InputToTensorConstantSpec tensor_constant;
+  50: InputToCustomObjSpec custom_obj;
+  70: InputTokenSpec token;
+  60: InputToConstantInputSpec constant_input;
+}
+
+struct UserOutputSpec {
+  10: Argument arg;
+}
+
+struct LossOutputSpec {
+  10: TensorArgument arg;
+}
+
+struct BufferMutationSpec {
+  10: TensorArgument arg;
+  20: string buffer_name;
+}
+
+struct ParameterMutationSpec {
+  10: TensorArgument arg;
+  20: string parameter_name;
+}
+
+struct GradientToParameterSpec {
+  10: TensorArgument arg;
+  20: string parameter_name;
+}
+
+struct GradientToUserInputSpec {
+  10: TensorArgument arg;
+  20: string user_input_name;
+}
+
+struct UserInputMutationSpec {
+  10: TensorArgument arg;
+  20: string user_input_name;
+}
+
+struct OutputTokenSpec {
+  10: TokenArgument arg;
+}
+
+union OutputSpec {
+  10: UserOutputSpec user_output;
+  20: LossOutputSpec loss_output;
+  30: BufferMutationSpec buffer_mutation;
+  40: GradientToParameterSpec gradient_to_parameter;
+  50: GradientToUserInputSpec gradient_to_user_input;
+  60: UserInputMutationSpec user_input_mutation;
+  70: OutputTokenSpec token;
+  80: ParameterMutationSpec parameter_mutation;
+}
+
+struct GraphSignature {
+  10: list<InputSpec> input_specs;
+  20: list<OutputSpec> output_specs;
+}
+
+struct RangeConstraint {
+  10: optional i64 min_val;
+  20: optional i64 max_val;
+}
+
+struct ModuleCallSignature {
+  10: list<Argument> inputs;
+  20: list<Argument> outputs;
+  30: string in_spec;
+  40: string out_spec;
+  50: optional list<string> forward_arg_names;
+}
+
+struct ModuleCallEntry {
+  10: string fqn;
+  30: optional ModuleCallSignature signature;
+}
+
+struct NamedTupleDef {
+  10: list<string> field_names;
+}
+
+struct GraphModule {
+  10: Graph graph;
+  50: GraphSignature signature;
+  60: list<ModuleCallEntry> module_call_graph;
+  40: map<string, string> metadata;
+  70: map<string, NamedTupleDef> treespec_namedtuple_fields;
+}
+
+struct SchemaVersion {
+  10: i64 major;
+  20: i64 minor;
+}
+
+struct ExportedProgram {
+  10: GraphModule graph_module;
+  20: map<string, i64> opset_version;
+  30: map<string, RangeConstraint> range_constraints;
+  60: SchemaVersion schema_version;
+  70: list<string> verifiers;
+  80: string torch_version;
+  90: list<string> guards_code;
+}
+
+struct PayloadMeta {
+  10: string path_name;
+  20: bool is_param;
+  30: bool use_pickle;
+  40: optional TensorMeta tensor_meta;
+}
+
+struct PayloadConfig {
+  10: map<string, PayloadMeta> config;
+}
+
+struct AOTInductorModelPickleData {
+  1: string library_basename;
+  2: list<string> input_names;
+  3: list<string> output_names;
+  4: optional i64 floating_point_input_dtype;
+  5: optional i64 floating_point_output_dtype;
+  6: optional bool aot_inductor_model_is_cpu;
+}
+
+struct ExternKernelNode {
+  10: string name;
+  20: Node node;
+}
+
+struct ExternKernelNodes {
+  10: list<ExternKernelNode> nodes;
+}

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/schema.py

@@ -0,0 +1,574 @@
+# NOTE: This is a placeholder for iterating on export serialization schema design.
+#       Anything is subject to change and no guarantee is provided at this point.
+
+from dataclasses import dataclass, field
+from enum import IntEnum
+from typing import Annotated
+
+from torch._export.serde.union import _Union, _union_dataclass
+
+
+# NOTE: Please update this value if any modifications are made to the schema
+SCHEMA_VERSION = (8, 20)
+TREESPEC_VERSION = 1
+
+
+# NOTE: If you updated the schema, please run `scripts/export/update_schema.py`
+# to update the auto generated files.
+#
+# There are also mappings from serialized enum values to c10 enum member names.
+# These are used by scripts/export/update_schema.py to generate C++ conversion
+# functions for oss_proxy_executor.cpp.
+# When adding new enum values, update both the enum AND the mapping.
+class ScalarType(IntEnum):
+    UNKNOWN = 0
+    BYTE = 1
+    CHAR = 2
+    SHORT = 3
+    INT = 4
+    LONG = 5
+    HALF = 6
+    FLOAT = 7
+    DOUBLE = 8
+    COMPLEXHALF = 9
+    COMPLEXFLOAT = 10
+    COMPLEXDOUBLE = 11
+    BOOL = 12
+    BFLOAT16 = 13
+    UINT16 = 28
+    FLOAT8E4M3FN = 29
+    FLOAT8E5M2 = 30
+    FLOAT8E4M3FNUZ = 31
+    FLOAT8E5M2FNUZ = 32
+    FLOAT8E8M0FNU = 33
+    UINT32 = 34
+    UINT64 = 35
+
+
+class Layout(IntEnum):
+    Unknown = 0
+    SparseCoo = 1
+    SparseCsr = 2
+    SparseCsc = 3
+    SparseBsr = 4
+    SparseBsc = 5
+    _mkldnn = 6
+    Strided = 7
+
+
+class MemoryFormat(IntEnum):
+    Unknown = 0
+    ContiguousFormat = 1
+    ChannelsLast = 2
+    ChannelsLast3d = 3
+    PreserveFormat = 4
+
+
+SCALAR_TYPE_TO_C10: dict[int, str] = {
+    ScalarType.BYTE: "Byte",
+    ScalarType.CHAR: "Char",
+    ScalarType.SHORT: "Short",
+    ScalarType.INT: "Int",
+    ScalarType.LONG: "Long",
+    ScalarType.HALF: "Half",
+    ScalarType.FLOAT: "Float",
+    ScalarType.DOUBLE: "Double",
+    ScalarType.COMPLEXHALF: "ComplexHalf",
+    ScalarType.COMPLEXFLOAT: "ComplexFloat",
+    ScalarType.COMPLEXDOUBLE: "ComplexDouble",
+    ScalarType.BOOL: "Bool",
+    ScalarType.BFLOAT16: "BFloat16",
+    ScalarType.UINT16: "UInt16",
+    ScalarType.FLOAT8E4M3FN: "Float8_e4m3fn",
+    ScalarType.FLOAT8E5M2: "Float8_e5m2",
+    ScalarType.FLOAT8E4M3FNUZ: "Float8_e4m3fnuz",
+    ScalarType.FLOAT8E5M2FNUZ: "Float8_e5m2fnuz",
+    ScalarType.FLOAT8E8M0FNU: "Float8_e8m0fnu",
+    ScalarType.UINT32: "UInt32",
+    ScalarType.UINT64: "UInt64",
+}
+
+LAYOUT_TO_C10: dict[int, str] = {
+    Layout.SparseCoo: "Sparse",
+    Layout.SparseCsr: "SparseCsr",
+    Layout.SparseCsc: "SparseCsc",
+    Layout.SparseBsr: "SparseBsr",
+    Layout.SparseBsc: "SparseBsc",
+    Layout._mkldnn: "Mkldnn",
+    Layout.Strided: "Strided",
+}
+
+MEMORY_FORMAT_TO_C10: dict[int, str] = {
+    MemoryFormat.ContiguousFormat: "Contiguous",
+    MemoryFormat.ChannelsLast: "ChannelsLast",
+    MemoryFormat.ChannelsLast3d: "ChannelsLast3d",
+    MemoryFormat.PreserveFormat: "Preserve",
+}
+
+
+@dataclass
+class Device:
+    type: Annotated[str, 10]
+    index: Annotated[int | None, 20] = None
+
+
+@_union_dataclass
+class SymExprHint(_Union):
+    as_int: Annotated[int, 10]
+    as_bool: Annotated[bool, 20]
+    as_float: Annotated[float, 30]
+
+
+# This is for storing the symbolic expressions behind symints/symfloats/symbools
+# For example, we can get something like
+# SymExpr(expr_str="s0 + s1", hint=SymExprHint(as_int=4)
+# if we also have the hint that s0 and s1 are both 2.
+@dataclass
+class SymExpr:
+    expr_str: Annotated[str, 10]
+    hint: Annotated[SymExprHint | None, 20] = None
+
+
+@_union_dataclass
+class SymInt(_Union):
+    as_expr: Annotated[SymExpr, 10]
+    as_int: Annotated[int, 20]
+
+
+@_union_dataclass
+class SymFloat(_Union):
+    as_expr: Annotated[SymExpr, 10]
+    as_float: Annotated[float, 20]
+
+
+@_union_dataclass
+class SymBool(_Union):
+    as_expr: Annotated[SymExpr, 10]
+    as_bool: Annotated[bool, 20]
+
+
+@dataclass
+class TensorMeta:
+    dtype: Annotated[ScalarType, 10]
+    sizes: Annotated[list[SymInt], 20]
+    requires_grad: Annotated[bool, 30]
+    device: Annotated[Device, 40]
+    strides: Annotated[list[SymInt], 50]
+    storage_offset: Annotated[SymInt, 60]
+    layout: Annotated[Layout, 70]
+
+
+# In most cases we will use the "as_name" field to store arguments which are
+# SymInts.
+# The "as_int" field is used in the case where we have a list containing a mix
+# of SymInt and ints (ex. [1, s0, ...]). We will serialize this type of list to
+# be List[SymIntArgument] and map the SymInts to the "as_name" field, and ints
+# to the "as_int" field.
+@_union_dataclass
+class SymIntArgument(_Union):
+    as_name: Annotated[str, 10]
+    as_int: Annotated[int, 20]
+
+
+# In most cases we will use the "as_name" field to store arguments which are
+# SymFloats.
+# The "as_float" field is used in the case where we have a list containing a mix
+# of SymFloat and float (ex. [1.0, s0, ...]). We will serialize this type of list to
+# be List[SymFloatArgument] and map the SymFloats to the "as_name" field, and ints
+# to the "as_float" field.
+@_union_dataclass
+class SymFloatArgument(_Union):
+    as_name: Annotated[str, 10]
+    as_float: Annotated[float, 20]
+
+
+# In most cases we will use the "as_name" field to store arguments which are
+# SymBools.
+# The "as_bool" field is used in the case where we have a list containing a mix
+# of SymBool and bools (ex. [True, i0, ...]). We will serialize this type of list to
+# be List[SymboolArgument] and map the SymBools to the "as_name" field, and bools
+# to the "as_bool" field.
+@_union_dataclass
+class SymBoolArgument(_Union):
+    as_name: Annotated[str, 10]
+    as_bool: Annotated[bool, 20]
+
+
+@dataclass
+class TensorArgument:
+    name: Annotated[str, 10]
+
+
+@dataclass
+class TokenArgument:
+    name: Annotated[str, 10]
+
+
+# This is use for storing the contents of a list which contain optional tensors
+# (Tensor?[], ex. [Tensor, None, ...]), where the list will be serialized to the
+# type List[OptionalTensorArgument], with tensor values serialized to the
+# "as_tensor" field, and None values serialized to the "as_none" field.
+@_union_dataclass
+class OptionalTensorArgument(_Union):
+    as_tensor: Annotated[TensorArgument, 20]
+    as_none: Annotated[bool, 10]
+
+
+@dataclass
+class GraphArgument:
+    name: Annotated[str, 10]
+    graph: Annotated["Graph", 20]
+
+
+@dataclass
+class CustomObjArgument:
+    name: Annotated[str, 10]
+    class_fqn: Annotated[str, 20]
+
+
+@dataclass
+class ComplexValue:
+    real: Annotated[float, 10]
+    imag: Annotated[float, 20]
+
+
+# This is actually a union type
+@_union_dataclass
+class Argument(_Union):
+    as_none: Annotated[bool, 10]
+    as_tensor: Annotated[TensorArgument, 20]
+    as_tensors: Annotated[list[TensorArgument], 30]
+    as_int: Annotated[int, 50]
+    as_ints: Annotated[list[int], 70]
+    as_float: Annotated[float, 80]
+    as_floats: Annotated[list[float], 90]
+    as_string: Annotated[str, 100]
+    as_strings: Annotated[list[str], 101]
+    as_sym_int: Annotated[SymIntArgument, 110]
+    as_sym_ints: Annotated[list[SymIntArgument], 120]
+    as_scalar_type: Annotated[ScalarType, 130]
+    as_memory_format: Annotated[MemoryFormat, 140]
+    as_layout: Annotated[Layout, 150]
+    as_device: Annotated[Device, 160]
+    as_bool: Annotated[bool, 170]
+    as_bools: Annotated[list[bool], 180]
+    as_sym_bool: Annotated[SymBoolArgument, 182]
+    as_sym_bools: Annotated[list[SymBoolArgument], 184]
+    as_graph: Annotated[GraphArgument, 200]
+    as_optional_tensors: Annotated[list[OptionalTensorArgument], 190]
+    as_custom_obj: Annotated[CustomObjArgument, 210]
+    as_operator: Annotated[str, 220]
+    as_sym_float: Annotated[SymFloatArgument, 230]
+    as_sym_floats: Annotated[list[SymFloatArgument], 240]
+    as_optional_tensor: Annotated[OptionalTensorArgument, 250]
+    as_complex: Annotated[ComplexValue, 260]
+    as_nested_tensors: Annotated[list[list[TensorArgument]], 270]
+    as_int_lists: Annotated[list[list[int]], 280]
+    as_string_to_argument: Annotated[dict[str, "Argument"], 290]
+    as_float_lists: Annotated[list[list[float]], 300]
+
+
+class ArgumentKind(IntEnum):
+    UNKNOWN = 0
+    POSITIONAL = 1
+    KEYWORD = 2
+
+
+@dataclass
+class NamedArgument:
+    # Argument name from the operator schema
+    name: Annotated[str, 10]
+    arg: Annotated[Argument, 20]
+    kind: Annotated[ArgumentKind | None, 30] = None
+
+
+@dataclass
+class Node:
+    target: Annotated[str, 10]
+    inputs: Annotated[list[NamedArgument], 20]
+    outputs: Annotated[list[Argument], 30]
+    metadata: Annotated[dict[str, str], 40]
+    is_hop_single_tensor_return: Annotated[bool | None, 50] = None
+    # For BC, default is None so older serialized models without 'name' can be loaded.
+    name: Annotated[str | None, 60] = None
+
+
+@dataclass
+class Graph:
+    inputs: Annotated[list[Argument], 10]
+    outputs: Annotated[list[Argument], 20]
+    nodes: Annotated[list[Node], 30]
+    tensor_values: Annotated[dict[str, TensorMeta], 40]
+    sym_int_values: Annotated[dict[str, SymInt], 50]
+    sym_bool_values: Annotated[dict[str, SymBool], 60]
+    # This is for deserializing the submodule graphs from higher order ops
+    # (ex. cond, map) where single tensor returns will just return a single
+    # tensor, rather than following export schema and returning a singleton
+    # list.
+    is_single_tensor_return: Annotated[bool, 70] = False
+    custom_obj_values: Annotated[dict[str, CustomObjArgument], 80] = field(
+        default_factory=dict
+    )
+    sym_float_values: Annotated[dict[str, SymFloat], 90] = field(default_factory=dict)
+
+
+@dataclass
+class UserInputSpec:
+    # Actually, only tensors and SymInts are allowed here
+    arg: Annotated[Argument, 10]
+
+
+@_union_dataclass
+class ConstantValue(_Union):
+    as_none: Annotated[bool, 10]
+    as_int: Annotated[int, 20]
+    as_float: Annotated[float, 30]
+    as_string: Annotated[str, 40]
+    as_bool: Annotated[bool, 50]
+
+
+@dataclass
+class InputToConstantInputSpec:
+    name: Annotated[str, 10]
+    value: Annotated[ConstantValue, 20]
+
+
+@dataclass
+class InputToParameterSpec:
+    arg: Annotated[TensorArgument, 10]
+    parameter_name: Annotated[str, 20]
+
+
+@dataclass
+class InputToBufferSpec:
+    arg: Annotated[TensorArgument, 10]
+    buffer_name: Annotated[str, 20]
+    persistent: Annotated[bool, 30]
+
+
+@dataclass
+class InputToTensorConstantSpec:
+    arg: Annotated[TensorArgument, 10]
+    tensor_constant_name: Annotated[str, 20]
+
+
+@dataclass
+class InputToCustomObjSpec:
+    arg: Annotated[CustomObjArgument, 10]
+    custom_obj_name: Annotated[str, 20]
+
+
+@dataclass
+class InputTokenSpec:
+    arg: Annotated[TokenArgument, 10]
+
+
+@_union_dataclass
+class InputSpec(_Union):
+    user_input: Annotated[UserInputSpec, 10]
+    parameter: Annotated[InputToParameterSpec, 20]
+    buffer: Annotated[InputToBufferSpec, 30]
+    tensor_constant: Annotated[InputToTensorConstantSpec, 40]
+    custom_obj: Annotated[InputToCustomObjSpec, 50]
+    token: Annotated[InputTokenSpec, 70]
+    constant_input: Annotated[InputToConstantInputSpec, 60]
+
+
+@dataclass
+class UserOutputSpec:
+    arg: Annotated[Argument, 10]
+
+
+@dataclass
+class LossOutputSpec:
+    arg: Annotated[TensorArgument, 10]
+
+
+@dataclass
+class BufferMutationSpec:
+    arg: Annotated[TensorArgument, 10]
+    buffer_name: Annotated[str, 20]
+
+
+@dataclass
+class ParameterMutationSpec:
+    arg: Annotated[TensorArgument, 10]
+    parameter_name: Annotated[str, 20]
+
+
+@dataclass
+class GradientToParameterSpec:
+    arg: Annotated[TensorArgument, 10]
+    parameter_name: Annotated[str, 20]
+
+
+@dataclass
+class GradientToUserInputSpec:
+    arg: Annotated[TensorArgument, 10]
+    user_input_name: Annotated[str, 20]
+
+
+@dataclass
+class UserInputMutationSpec:
+    arg: Annotated[TensorArgument, 10]
+    user_input_name: Annotated[str, 20]
+
+
+@dataclass
+class OutputTokenSpec:
+    arg: Annotated[TokenArgument, 10]
+
+
+@_union_dataclass
+class OutputSpec(_Union):
+    user_output: Annotated[UserOutputSpec, 10]
+    loss_output: Annotated[LossOutputSpec, 20]
+    buffer_mutation: Annotated[BufferMutationSpec, 30]
+    gradient_to_parameter: Annotated[GradientToParameterSpec, 40]
+    gradient_to_user_input: Annotated[GradientToUserInputSpec, 50]
+    user_input_mutation: Annotated[UserInputMutationSpec, 60]
+    token: Annotated[OutputTokenSpec, 70]
+    parameter_mutation: Annotated[ParameterMutationSpec, 80]
+
+
+@dataclass
+class GraphSignature:
+    input_specs: Annotated[list[InputSpec], 10]
+    output_specs: Annotated[list[OutputSpec], 20]
+
+
+@dataclass
+class RangeConstraint:
+    min_val: Annotated[int | None, 10]
+    max_val: Annotated[int | None, 20]
+
+
+@dataclass
+class ModuleCallSignature:
+    inputs: Annotated[list[Argument], 10]
+    outputs: Annotated[list[Argument], 20]
+
+    # These are serialized by calling pytree.treespec_loads
+    # And deserialized by calling pytree.treespec_dumps
+    in_spec: Annotated[str, 30]
+    out_spec: Annotated[str, 40]
+
+    # This field is used to prettify the graph placeholders
+    # after we Ser/Der and retrace
+    forward_arg_names: Annotated[list[str] | None, 50] = None
+
+
+@dataclass
+class ModuleCallEntry:
+    fqn: Annotated[str, 10]
+    signature: Annotated[ModuleCallSignature | None, 30] = None
+
+
+@dataclass
+class NamedTupleDef:
+    field_names: Annotated[list[str], 10]
+
+
+@dataclass
+class GraphModule:
+    graph: Annotated[Graph, 10]
+    signature: Annotated[GraphSignature, 50]
+    # This is used for unflattening, by tracking the calling structure of all of
+    # the modules in order to unflatten the modules back to the eager calling
+    # conventions.
+    module_call_graph: Annotated[list[ModuleCallEntry], 60]
+    metadata: Annotated[dict[str, str], 40] = field(default_factory=dict)
+    # Mapping of namedtuple types to namedtuple field names, used for BC
+    treespec_namedtuple_fields: Annotated[dict[str, NamedTupleDef], 70] = field(
+        default_factory=dict
+    )
+
+
+# Invariant: Every time a change is made to the schema, one of the versions
+#            should be updated.
+@dataclass
+class SchemaVersion:
+    major: Annotated[
+        int, 10
+    ]  # Major version number is bumped every time a breaking change is made.
+    minor: Annotated[
+        int, 20
+    ]  # Minor version number is bumped when a compatible change is made.
+
+
+@dataclass
+class ExportedProgram:
+    graph_module: Annotated[GraphModule, 10]
+    # Key is the opset namespace (ex. aten), and value is the version number
+    opset_version: Annotated[dict[str, int], 20]
+    range_constraints: Annotated[dict[str, RangeConstraint], 30]
+    schema_version: Annotated[SchemaVersion, 60]
+    verifiers: Annotated[list[str], 70] = field(default_factory=list)
+    torch_version: Annotated[str, 80] = "<=2.4"
+    guards_code: Annotated[list[str], 90] = field(default_factory=list)
+
+
+#########################################################################
+# Container types for inference tasks, not being used directly for export.
+#########################################################################
+
+
+# The metadata for payload saved in PT2 archive.
+# payload includes params, buffers, tensor constants, and custom objects.
+@dataclass
+class PayloadMeta:
+    # the path of the payload in the archive file, e.g. "weight_0"
+    path_name: Annotated[str, 10]
+    is_param: Annotated[bool, 20]
+    # whether the payload is serialized using pickle.
+    # Only custom objects and tensor subclasses that are not fake tensors
+    # are serialized using pickle.
+    use_pickle: Annotated[bool, 30]
+    # Custom Objects don't have tensor_meta and will be serialized using pickle
+    tensor_meta: Annotated[TensorMeta | None, 40]
+
+
+# The mapping from payload FQN to its metadata.
+@dataclass
+class PayloadConfig:
+    config: Annotated[dict[str, PayloadMeta], 10]
+
+
+#
+# The structure is used to serialize instances of AOTInductorModel to pass
+# them from the publishing pipeline to the predictor.
+#
+# All new fields should be marked as optional.
+#
+@dataclass
+class AOTInductorModelPickleData:
+    # Base name of an associated .so AOTInductor library. Typically looks like:
+    # "abc.so".
+    library_basename: Annotated[str, 1]
+
+    # AOTInductor engine input names.
+    input_names: Annotated[list[str], 2]
+
+    # AOTInductor engine output names.
+    output_names: Annotated[list[str], 3]
+
+    # These fields tell whether floating point inputs/outputs should be converted to
+    # a certain type. If None, the dtypes that the AOTInductor engine inferred from the sample
+    # inputs are used.
+    floating_point_input_dtype: Annotated[int | None, 4] = None
+    floating_point_output_dtype: Annotated[int | None, 5] = None
+
+    # Whether AOTInductor runtime is for CPU.
+    aot_inductor_model_is_cpu: Annotated[bool | None, 6] = None
+
+
+@dataclass
+class ExternKernelNode:
+    # name is not the unique identifier of the node
+    name: Annotated[str, 10]
+    node: Annotated[Node, 20]
+
+
+@dataclass
+class ExternKernelNodes:
+    nodes: Annotated[list[ExternKernelNode], 10]

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/schema.yaml

@@ -0,0 +1,569 @@
+# @generated by update_schema.py
+# checksum<<eb9645ae5f539caace53ada70fb262097ce780e6edf685149dafc5a5abe8d31b>>
+AOTInductorModelPickleData:
+  kind: struct
+  fields:
+    library_basename:
+      type: str
+    input_names:
+      type: List[str]
+    output_names:
+      type: List[str]
+    floating_point_input_dtype:
+      type: Optional[int]
+      default: None
+    floating_point_output_dtype:
+      type: Optional[int]
+      default: None
+    aot_inductor_model_is_cpu:
+      type: Optional[bool]
+      default: None
+Argument:
+  kind: union
+  fields:
+    as_none:
+      type: bool
+    as_tensor:
+      type: TensorArgument
+    as_tensors:
+      type: List[TensorArgument]
+    as_int:
+      type: int
+    as_ints:
+      type: List[int]
+    as_float:
+      type: float
+    as_floats:
+      type: List[float]
+    as_string:
+      type: str
+    as_strings:
+      type: List[str]
+    as_sym_int:
+      type: SymIntArgument
+    as_sym_ints:
+      type: List[SymIntArgument]
+    as_scalar_type:
+      type: ScalarType
+    as_memory_format:
+      type: MemoryFormat
+    as_layout:
+      type: Layout
+    as_device:
+      type: Device
+    as_bool:
+      type: bool
+    as_bools:
+      type: List[bool]
+    as_sym_bool:
+      type: SymBoolArgument
+    as_sym_bools:
+      type: List[SymBoolArgument]
+    as_graph:
+      type: GraphArgument
+    as_optional_tensors:
+      type: List[OptionalTensorArgument]
+    as_custom_obj:
+      type: CustomObjArgument
+    as_operator:
+      type: str
+    as_sym_float:
+      type: SymFloatArgument
+    as_sym_floats:
+      type: List[SymFloatArgument]
+    as_optional_tensor:
+      type: OptionalTensorArgument
+    as_complex:
+      type: ComplexValue
+    as_nested_tensors:
+      type: List[List[TensorArgument]]
+    as_int_lists:
+      type: List[List[int]]
+    as_string_to_argument:
+      type: Dict[str, Argument]
+    as_float_lists:
+      type: List[List[float]]
+ArgumentKind:
+  kind: enum
+  fields:
+    UNKNOWN: 0
+    POSITIONAL: 1
+    KEYWORD: 2
+BufferMutationSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    buffer_name:
+      type: str
+ComplexValue:
+  kind: struct
+  fields:
+    real:
+      type: float
+    imag:
+      type: float
+ConstantValue:
+  kind: union
+  fields:
+    as_none:
+      type: bool
+    as_int:
+      type: int
+    as_float:
+      type: float
+    as_string:
+      type: str
+    as_bool:
+      type: bool
+CustomObjArgument:
+  kind: struct
+  fields:
+    name:
+      type: str
+    class_fqn:
+      type: str
+Device:
+  kind: struct
+  fields:
+    type:
+      type: str
+    index:
+      type: Optional[int]
+      default: None
+ExportedProgram:
+  kind: struct
+  fields:
+    graph_module:
+      type: GraphModule
+    opset_version:
+      type: Dict[str, int]
+    range_constraints:
+      type: Dict[str, RangeConstraint]
+    schema_version:
+      type: SchemaVersion
+    verifiers:
+      type: List[str]
+      default: '[]'
+    torch_version:
+      type: str
+      default: <=2.4
+    guards_code:
+      type: List[str]
+      default: '[]'
+ExternKernelNode:
+  kind: struct
+  fields:
+    name:
+      type: str
+    node:
+      type: Node
+ExternKernelNodes:
+  kind: struct
+  fields:
+    nodes:
+      type: List[ExternKernelNode]
+GradientToParameterSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    parameter_name:
+      type: str
+GradientToUserInputSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    user_input_name:
+      type: str
+Graph:
+  kind: struct
+  fields:
+    inputs:
+      type: List[Argument]
+    outputs:
+      type: List[Argument]
+    nodes:
+      type: List[Node]
+    tensor_values:
+      type: Dict[str, TensorMeta]
+    sym_int_values:
+      type: Dict[str, SymInt]
+    sym_bool_values:
+      type: Dict[str, SymBool]
+    is_single_tensor_return:
+      type: bool
+      default: 'False'
+    custom_obj_values:
+      type: Dict[str, CustomObjArgument]
+      default: '{}'
+    sym_float_values:
+      type: Dict[str, SymFloat]
+      default: '{}'
+GraphArgument:
+  kind: struct
+  fields:
+    name:
+      type: str
+    graph:
+      type: Graph
+GraphModule:
+  kind: struct
+  fields:
+    graph:
+      type: Graph
+    signature:
+      type: GraphSignature
+    module_call_graph:
+      type: List[ModuleCallEntry]
+    metadata:
+      type: Dict[str, str]
+      default: '{}'
+    treespec_namedtuple_fields:
+      type: Dict[str, NamedTupleDef]
+      default: '{}'
+GraphSignature:
+  kind: struct
+  fields:
+    input_specs:
+      type: List[InputSpec]
+    output_specs:
+      type: List[OutputSpec]
+InputSpec:
+  kind: union
+  fields:
+    user_input:
+      type: UserInputSpec
+    parameter:
+      type: InputToParameterSpec
+    buffer:
+      type: InputToBufferSpec
+    tensor_constant:
+      type: InputToTensorConstantSpec
+    custom_obj:
+      type: InputToCustomObjSpec
+    token:
+      type: InputTokenSpec
+    constant_input:
+      type: InputToConstantInputSpec
+InputToBufferSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    buffer_name:
+      type: str
+    persistent:
+      type: bool
+InputToConstantInputSpec:
+  kind: struct
+  fields:
+    name:
+      type: str
+    value:
+      type: ConstantValue
+InputToCustomObjSpec:
+  kind: struct
+  fields:
+    arg:
+      type: CustomObjArgument
+    custom_obj_name:
+      type: str
+InputToParameterSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    parameter_name:
+      type: str
+InputToTensorConstantSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    tensor_constant_name:
+      type: str
+InputTokenSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TokenArgument
+Layout:
+  kind: enum
+  fields:
+    Unknown: 0
+    SparseCoo: 1
+    SparseCsr: 2
+    SparseCsc: 3
+    SparseBsr: 4
+    SparseBsc: 5
+    _mkldnn: 6
+    Strided: 7
+LossOutputSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+MemoryFormat:
+  kind: enum
+  fields:
+    Unknown: 0
+    ContiguousFormat: 1
+    ChannelsLast: 2
+    ChannelsLast3d: 3
+    PreserveFormat: 4
+ModuleCallEntry:
+  kind: struct
+  fields:
+    fqn:
+      type: str
+    signature:
+      type: Optional[ModuleCallSignature]
+      default: None
+ModuleCallSignature:
+  kind: struct
+  fields:
+    inputs:
+      type: List[Argument]
+    outputs:
+      type: List[Argument]
+    in_spec:
+      type: str
+    out_spec:
+      type: str
+    forward_arg_names:
+      type: Optional[List[str]]
+      default: None
+NamedArgument:
+  kind: struct
+  fields:
+    name:
+      type: str
+    arg:
+      type: Argument
+    kind:
+      type: Optional[ArgumentKind]
+      default: None
+NamedTupleDef:
+  kind: struct
+  fields:
+    field_names:
+      type: List[str]
+Node:
+  kind: struct
+  fields:
+    target:
+      type: str
+    inputs:
+      type: List[NamedArgument]
+    outputs:
+      type: List[Argument]
+    metadata:
+      type: Dict[str, str]
+    is_hop_single_tensor_return:
+      type: Optional[bool]
+      default: None
+    name:
+      type: Optional[str]
+      default: None
+OptionalTensorArgument:
+  kind: union
+  fields:
+    as_tensor:
+      type: TensorArgument
+    as_none:
+      type: bool
+OutputSpec:
+  kind: union
+  fields:
+    user_output:
+      type: UserOutputSpec
+    loss_output:
+      type: LossOutputSpec
+    buffer_mutation:
+      type: BufferMutationSpec
+    gradient_to_parameter:
+      type: GradientToParameterSpec
+    gradient_to_user_input:
+      type: GradientToUserInputSpec
+    user_input_mutation:
+      type: UserInputMutationSpec
+    token:
+      type: OutputTokenSpec
+    parameter_mutation:
+      type: ParameterMutationSpec
+OutputTokenSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TokenArgument
+ParameterMutationSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    parameter_name:
+      type: str
+PayloadConfig:
+  kind: struct
+  fields:
+    config:
+      type: Dict[str, PayloadMeta]
+PayloadMeta:
+  kind: struct
+  fields:
+    path_name:
+      type: str
+    is_param:
+      type: bool
+    use_pickle:
+      type: bool
+    tensor_meta:
+      type: Optional[TensorMeta]
+RangeConstraint:
+  kind: struct
+  fields:
+    min_val:
+      type: Optional[int]
+    max_val:
+      type: Optional[int]
+ScalarType:
+  kind: enum
+  fields:
+    UNKNOWN: 0
+    BYTE: 1
+    CHAR: 2
+    SHORT: 3
+    INT: 4
+    LONG: 5
+    HALF: 6
+    FLOAT: 7
+    DOUBLE: 8
+    COMPLEXHALF: 9
+    COMPLEXFLOAT: 10
+    COMPLEXDOUBLE: 11
+    BOOL: 12
+    BFLOAT16: 13
+    UINT16: 28
+    FLOAT8E4M3FN: 29
+    FLOAT8E5M2: 30
+    FLOAT8E4M3FNUZ: 31
+    FLOAT8E5M2FNUZ: 32
+    FLOAT8E8M0FNU: 33
+    UINT32: 34
+    UINT64: 35
+SchemaVersion:
+  kind: struct
+  fields:
+    major:
+      type: int
+    minor:
+      type: int
+SymBool:
+  kind: union
+  fields:
+    as_expr:
+      type: SymExpr
+    as_bool:
+      type: bool
+SymBoolArgument:
+  kind: union
+  fields:
+    as_name:
+      type: str
+    as_bool:
+      type: bool
+SymExpr:
+  kind: struct
+  fields:
+    expr_str:
+      type: str
+    hint:
+      type: Optional[SymExprHint]
+      default: None
+SymExprHint:
+  kind: union
+  fields:
+    as_int:
+      type: int
+    as_bool:
+      type: bool
+    as_float:
+      type: float
+SymFloat:
+  kind: union
+  fields:
+    as_expr:
+      type: SymExpr
+    as_float:
+      type: float
+SymFloatArgument:
+  kind: union
+  fields:
+    as_name:
+      type: str
+    as_float:
+      type: float
+SymInt:
+  kind: union
+  fields:
+    as_expr:
+      type: SymExpr
+    as_int:
+      type: int
+SymIntArgument:
+  kind: union
+  fields:
+    as_name:
+      type: str
+    as_int:
+      type: int
+TensorArgument:
+  kind: struct
+  fields:
+    name:
+      type: str
+TensorMeta:
+  kind: struct
+  fields:
+    dtype:
+      type: ScalarType
+    sizes:
+      type: List[SymInt]
+    requires_grad:
+      type: bool
+    device:
+      type: Device
+    strides:
+      type: List[SymInt]
+    storage_offset:
+      type: SymInt
+    layout:
+      type: Layout
+TokenArgument:
+  kind: struct
+  fields:
+    name:
+      type: str
+UserInputMutationSpec:
+  kind: struct
+  fields:
+    arg:
+      type: TensorArgument
+    user_input_name:
+      type: str
+UserInputSpec:
+  kind: struct
+  fields:
+    arg:
+      type: Argument
+UserOutputSpec:
+  kind: struct
+  fields:
+    arg:
+      type: Argument
+SCHEMA_VERSION:
+- 8
+- 20
+TREESPEC_VERSION: 1

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/schema_check.py

@@ -0,0 +1,905 @@
+# mypy: allow-untyped-defs
+import dataclasses
+import hashlib
+import inspect
+import re
+import types
+import typing
+from enum import IntEnum
+from typing import Annotated, Any, ForwardRef, Union
+
+from torch._export.serde import schema
+from torch._export.serde.union import _Union
+
+
+class SchemaUpdateError(Exception):
+    pass
+
+
+def _check(x, msg):
+    if not x:
+        raise SchemaUpdateError(msg)
+
+
+_CPP_TYPE_MAP = {
+    str: "std::string",
+    int: "int64_t",
+    float: "F64",
+    bool: "bool",
+}
+
+_THRIFT_TYPE_MAP = {
+    str: "string",
+    int: "i64",
+    float: "double",
+    bool: "bool",
+}
+
+
+def _staged_schema():
+    yaml_ret: dict[str, Any] = {}
+    defs = {}
+    cpp_enum_defs: dict[str, str] = {}
+    cpp_class_defs: dict[str, str] = {}
+    cpp_type_decls: list[str] = []
+    cpp_json_defs: list[str] = []
+    thrift_enum_defs: list[str] = []
+    thrift_type_defs: dict[str, str] = {}
+
+    def _handle_aggregate(ty) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+        def dump_type(t, level: int) -> tuple[str, str, str]:
+            if getattr(t, "__name__", None) in cpp_enum_defs:
+                return t.__name__, "int64_t", t.__name__
+            elif t in _CPP_TYPE_MAP:
+                return (t.__name__, _CPP_TYPE_MAP[t], _THRIFT_TYPE_MAP[t])
+            elif isinstance(t, str):
+                if t not in defs:
+                    raise AssertionError(f"type {t} not in defs")
+                if t in cpp_enum_defs:
+                    raise AssertionError(f"type {t} unexpectedly in cpp_enum_defs")
+                if "[" in t:
+                    raise AssertionError(f"type {t} contains '[' which is not allowed")
+                return t, f"ForwardRef<{t}>", t
+            elif isinstance(t, ForwardRef):
+                return (
+                    t.__forward_arg__,
+                    f"ForwardRef<{t.__forward_arg__}>",
+                    t.__forward_arg__,
+                )
+            elif o := typing.get_origin(t):
+                # Lemme know if there's a better way to do this.
+                if o is list:
+                    yaml_head, cpp_head, thrift_head, thrift_tail = (
+                        "List",
+                        "std::vector",
+                        "list<",
+                        ">",
+                    )
+                elif o is dict:
+                    yaml_head, cpp_head, thrift_head, thrift_tail = (
+                        "Dict",
+                        "std::unordered_map",
+                        "map<",
+                        ">",
+                    )
+                elif o is Union or o is types.UnionType:
+                    if level != 0:
+                        raise AssertionError(
+                            f"Optional is only supported at the top level, got level={level}"
+                        )
+                    args = typing.get_args(t)
+                    if len(args) != 2 or args[1] is not type(None):
+                        raise AssertionError(
+                            f"expected Optional type with 2 args ending in None, got {args}"
+                        )
+                    yaml_type, cpp_type, thrift_type = dump_type(args[0], level + 1)
+                    return (
+                        f"Optional[{yaml_type}]",
+                        f"std::optional<{cpp_type}>",
+                        f"optional {thrift_type}",
+                    )
+                elif o is Annotated:
+                    return dump_type(t.__origin__, level)
+                else:
+                    raise AssertionError(f"Type {t} is not supported in export schema.")
+                yaml_arg_types, cpp_arg_types, thrift_arg_types = zip(
+                    *[dump_type(x, level + 1) for x in typing.get_args(t)]
+                )
+                return (
+                    (f"{yaml_head}[{', '.join(yaml_arg_types)}]"),
+                    (f"{cpp_head}<{', '.join(cpp_arg_types)}>"),
+                    f"{thrift_head}{', '.join(thrift_arg_types)}{thrift_tail}",
+                )
+            elif isinstance(t, type):
+                return (t.__name__, t.__name__, t.__name__)
+            else:
+                raise AssertionError(f"Type {t} is not supported in export schema.")
+
+        def dump_cpp_value(v) -> str:
+            if v is None:
+                return "std::nullopt"
+            elif v is True:
+                return "true"
+            elif v is False:
+                return "false"
+            elif v == {}:
+                return "{}"
+            elif v == []:
+                return "{}"
+            elif v == ():
+                return "{}"
+            elif isinstance(v, str):
+                return f'"{v}"'
+            else:
+                raise AssertionError(
+                    f"Default value {v} is not supported yet in export schema."
+                )
+
+        def dump_field(f) -> tuple[dict[str, Any], str, str | None, str, int]:
+            t, cpp_type, thrift_type = dump_type(f.type, 0)
+            ret = {"type": t}
+            cpp_default: str | None = None
+            if typing.get_origin(f.type) is not Annotated:
+                raise AssertionError(
+                    f"Field {f.name} must be annotated with an integer id."
+                )
+            thrift_id = f.type.__metadata__[0]
+            if type(thrift_id) is not int:
+                raise AssertionError(
+                    f"Field {f.name} must be annotated with an integer id, got {type(thrift_id)}"
+                )
+
+            value = dataclasses.MISSING
+            if f.default is not dataclasses.MISSING:
+                value = f.default
+            elif f.default_factory is not dataclasses.MISSING:
+                value = f.default_factory()
+
+            if value is not dataclasses.MISSING:
+                default = str(value)
+                ret["default"] = default
+                cpp_default = dump_cpp_value(value)
+
+                if t.startswith("Optional[") and value is not None:
+                    raise AssertionError(
+                        f"Optional field {ty.__name__}.{f.name} must have default value to be None."
+                    )
+
+            return ret, cpp_type, cpp_default, thrift_type, thrift_id
+
+        yaml_ret = {}
+        cpp_ret = {}
+        thrift_ret = {}
+        thrift_ids = set()
+        for f in dataclasses.fields(ty):
+            yaml_res, cpp_type, cpp_default, thrift_type, thrift_id = dump_field(f)
+            yaml_ret[f.name] = yaml_res
+            cpp_ret[f.name] = {"cpp_type": cpp_type, "cpp_default": cpp_default}
+            thrift_ret[f.name] = {"thrift_type": thrift_type, "thrift_id": thrift_id}
+            if thrift_id in thrift_ids:
+                raise AssertionError(
+                    f"Duplicate thrift id {thrift_id} for field {f.name} in {ty.__name__}."
+                )
+            thrift_ids.add(thrift_id)
+        return yaml_ret, cpp_ret, thrift_ret
+
+    def _handle_int_enum(name, ty):
+        yaml_ret[name] = {"kind": "enum", "fields": {x.name: x.value for x in ty}}
+        cpp_enum_defs[name] = f"""
+enum class {name} {{
+{chr(10).join([f"  {x.name} = {x.value}," for x in ty])}
+}};
+
+inline std::string_view printEnum(const {name}& e) {{
+  switch (e) {{
+{chr(10).join([f"    case {name}::{x.name}: return {chr(34)}{x.name}{chr(34)};" for x in ty])}
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }}
+}}
+
+inline void parseEnum(std::string_view s, {name}& t) {{
+{chr(10).join([f"  if (s == {chr(34)}{x.name}{chr(34)}) {{ t = {name}::{x.name}; return; }}" for x in ty])}
+  throw std::runtime_error("Unknown enum value: " + std::string{{s}});
+}}
+"""
+        thrift_enum_defs.append(
+            f"""
+enum {name} {{
+{chr(10).join([f"  {x.name} = {x.value}," for x in ty])}
+}}
+"""
+        )
+
+    def _handle_struct(name, ty):
+        fields, cpp_fields, thrift_fields = _handle_aggregate(ty)
+        yaml_ret[name] = {"kind": "struct", "fields": fields}
+        field_decls = "\n".join(
+            f"  {f['cpp_type']} {name}{' = ' + f['cpp_default'] if f['cpp_default'] is not None else ''};"
+            for name, f in cpp_fields.items()
+        )
+
+        def accessor(name, ty):
+            type_name = fields[name]["type"]
+            if type_name in cpp_enum_defs:
+                return f"""
+  {type_name} get_{name}() const {{
+    return static_cast<{type_name}>({name});
+  }}
+
+  void set_{name}({type_name} def) {{
+    {name} = static_cast<int64_t>(def);
+  }}
+"""
+            return f"""
+  const {ty}& get_{name}() const {{
+    return {name};
+  }}
+
+  void set_{name}({ty} def) {{
+    {name} = std::move(def);
+  }}
+"""
+
+        to_json_decl = f"void to_json(nlohmann::json& nlohmann_json_j, const {name}& nlohmann_json_t)"
+        to_json_def = f"""{{
+{chr(10).join([f'  nlohmann_json_j["{name}"] = nlohmann_json_t.{name};' for name, f in cpp_fields.items()])}
+}}
+"""
+        from_json_decl = f"void from_json(const nlohmann::json& nlohmann_json_j, {name}& nlohmann_json_t)"
+
+        from_json_def = f"""{{
+  {name} nlohmann_json_default_obj;
+{
+            chr(10).join(
+                [
+                    f'  nlohmann_json_t.{name} = nlohmann_json_j.value("{name}", nlohmann_json_default_obj.{name});'
+                    for name, f in cpp_fields.items()
+                ]
+            )
+        }
+}}
+"""
+        cpp_class_defs[name] = f"""
+class {name} {{
+ private:
+{field_decls}
+
+ public:
+{"".join([accessor(name, f["cpp_type"]) for name, f in cpp_fields.items()])}
+  friend {to_json_decl};
+  friend {from_json_decl};
+}};
+"""
+        cpp_json_defs.append(f"inline {to_json_decl} {to_json_def}")
+        cpp_json_defs.append(f"inline {from_json_decl} {from_json_def}")
+        cpp_type_decls.append(f"class {name};")
+
+        thrift_type_defs[name] = f"""
+struct {name} {{
+{chr(10).join(f"  {f['thrift_id']}: {f['thrift_type']} {n};" for n, f in thrift_fields.items())}
+}}"""
+
+    def _handle_union(name, ty):
+        fields, cpp_fields, thrift_fields = _handle_aggregate(ty)
+        yaml_ret[name] = {"kind": "union", "fields": fields}
+
+        def accessor(name, ty, idx):
+            return f"""
+  const {ty}& get_{name}() const {{
+    return std::get<{idx + 1}>(variant_);
+  }}
+
+  void set_{name}({ty} def) {{
+    variant_.emplace<{idx + 1}>(std::move(def));
+    tag_ = Tag::{name.upper()};
+  }}
+"""
+
+        to_json_branches = "".join(
+            [
+                f"""
+    if (nlohmann_json_t.tag_ == Tag::{name.upper()}) {{
+      nlohmann_json_j["{name}"] = nlohmann_json_t.get_{name}();
+      return;
+    }}"""
+                for idx, (name, f) in enumerate(cpp_fields.items())
+            ]
+        )
+        from_json_branches = "".join(
+            [
+                f"""
+    if (nlohmann_json_j.contains("{name}")) {{
+      nlohmann_json_t.variant_.emplace<{idx + 1}>(nlohmann_json_j.at("{name}").template get<{f["cpp_type"]}>());
+      nlohmann_json_t.tag_ = Tag::{name.upper()};
+      return;
+    }}"""
+                for idx, (name, f) in enumerate(cpp_fields.items())
+            ]
+        )
+
+        cpp_class_defs[name] = f"""
+class {name} {{
+  struct Void {{}};
+
+ public:
+  enum class Tag {{
+    {", ".join([name.upper() for name in cpp_fields])}
+  }};
+
+ private:
+  std::variant<Void, {", ".join(f["cpp_type"] for f in cpp_fields.values())}> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {{
+    return tag_;
+  }}
+{"".join([accessor(name, f["cpp_type"], idx) for idx, (name, f) in enumerate(cpp_fields.items())])}
+  friend void to_json(nlohmann::json& nlohmann_json_j, const {name}& nlohmann_json_t) {{
+{to_json_branches}
+  }}
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, {name}& nlohmann_json_t) {{
+{from_json_branches}
+  }}
+}};
+
+inline std::string_view printEnum(const {name}::Tag& e) {{
+  switch (e) {{
+{chr(10).join([f"    case {name}::Tag::{x.upper()}: return {chr(34)}{x.upper()}{chr(34)};" for x in cpp_fields])}
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }}
+}}
+
+inline void parseEnum(std::string_view s, {name}::Tag& t) {{
+{chr(10).join([f"  if (s == {chr(34)}{x.upper()}{chr(34)}) {{ t = {name}::Tag::{x.upper()}; return; }}" for x in cpp_fields])}
+  throw std::runtime_error("Unknown enum value: " + std::string{{s}});
+}}
+
+"""
+        cpp_type_decls.append(f"class {name};")
+
+        thrift_type_defs[name] = f"""
+union {name} {{
+{chr(10).join(f"  {f['thrift_id']}: {f['thrift_type']} {n};" for n, f in thrift_fields.items())}
+}}"""
+
+    for name in dir(schema):
+        if name.startswith("_"):
+            continue
+
+        value = getattr(schema, name)
+
+        if hasattr(value, "__module__") and value.__module__ != schema.__name__:
+            continue
+
+        defs[name] = value
+
+    class_ordering = {}
+    for name, value in defs.items():
+        if isinstance(value, type):
+            if issubclass(value, IntEnum):
+                _handle_int_enum(name, value)
+            elif dataclasses.is_dataclass(value):
+                class_ordering[name] = inspect.findsource(value)[1]
+                if issubclass(value, _Union):
+                    _handle_union(name, value)
+                else:
+                    _handle_struct(name, value)
+            else:
+                raise AssertionError(f"Unknown schema type {name}: {value}")
+        elif isinstance(value, (int, tuple)):
+            if name not in ("SCHEMA_VERSION", "TREESPEC_VERSION"):
+                raise AssertionError(
+                    f"expected SCHEMA_VERSION or TREESPEC_VERSION, got {name}"
+                )
+        elif isinstance(value, dict):
+            # Skip mapping dictionaries used for codegen
+            pass
+        else:
+            raise AssertionError(f"Unknown variable {name}: {value}")
+
+    yaml_ret["SCHEMA_VERSION"] = list(defs["SCHEMA_VERSION"])
+    if not all(x > 0 for x in yaml_ret["SCHEMA_VERSION"]):
+        raise AssertionError(
+            f"all SCHEMA_VERSION values must be > 0, got {yaml_ret['SCHEMA_VERSION']}"
+        )
+    yaml_ret["TREESPEC_VERSION"] = defs["TREESPEC_VERSION"]
+    if yaml_ret["TREESPEC_VERSION"] <= 0:
+        raise AssertionError(
+            f"TREESPEC_VERSION must be > 0, got {yaml_ret['TREESPEC_VERSION']}"
+        )
+
+    cpp_header = f"""
+#pragma once
+
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+
+#ifndef NLOHMANN_JSON_NAMESPACE_BEGIN
+#define NLOHMANN_JSON_NAMESPACE_BEGIN namespace nlohmann {{
+#endif
+
+#ifndef NLOHMANN_JSON_NAMESPACE_END
+#define NLOHMANN_JSON_NAMESPACE_END }}
+#endif
+
+// https://github.com/nlohmann/json/pull/2117
+NLOHMANN_JSON_NAMESPACE_BEGIN
+template <typename T>
+struct adl_serializer<std::optional<T>> {{
+  static void to_json(json& j, const std::optional<T>& opt) {{
+    if (opt == std::nullopt) {{
+      j = nullptr;
+    }} else {{
+      j = *opt; // this will call adl_serializer<T>::to_json which will
+                // find the free function to_json in T's namespace!
+    }}
+  }}
+
+  static void from_json(const json& j, std::optional<T>& opt) {{
+    if (j.is_null()) {{
+      opt = std::nullopt;
+    }} else {{
+      opt = j.template get<T>(); // same as above, but with
+                                 // adl_serializer<T>::from_json
+    }}
+  }}
+}};
+NLOHMANN_JSON_NAMESPACE_END
+
+namespace torch {{
+namespace _export {{
+
+template <typename T>
+class ForwardRef {{
+  static_assert(!std::is_reference_v<T>, "ForwardRef cannot be a reference type");
+
+ public:
+  ForwardRef(): ptr_(std::make_unique<T>()) {{}}
+  ForwardRef(ForwardRef<T>&&);
+  ForwardRef(const ForwardRef<T>& other): ptr_(std::make_unique<T>(*other.ptr_)) {{}}
+  ForwardRef<T>& operator=(ForwardRef<T>&&);
+  ForwardRef<T>& operator=(const ForwardRef<T>& other) {{
+    ptr_ = std::make_unique<T>(*other.ptr_);
+    return *this;
+  }}
+  ~ForwardRef();
+  const T& operator*() const {{
+    return *ptr_;
+  }}
+
+  const T* operator->() const {{
+    return ptr_.get();
+  }}
+
+  void emplace(T&& t) {{
+    ptr_ = std::make_unique<T>(std::move(t));
+  }}
+
+ private:
+  std::unique_ptr<T> ptr_;
+}};
+
+template <typename T>
+void to_json(nlohmann::json& j, const ForwardRef<T>& p) {{
+  j = *p;
+}}
+
+template <typename T>
+void from_json(const nlohmann::json& j, ForwardRef<T>& p) {{
+  p.emplace(j.template get<T>());
+}}
+
+class F64 {{
+ public:
+  double get() const {{
+    return value_;
+  }}
+
+  void set(double value) {{
+    value_ = value;
+  }}
+
+ private:
+  double value_;
+}};
+
+inline void to_json(nlohmann::json& j, const F64& f) {{
+  if (std::isinf(f.get())) {{
+    j = "Infinity";
+  }} else if (std::isinf(-f.get())) {{
+    j = "-Infinity";
+  }} else if (std::isnan(f.get())) {{
+    j = "NaN";
+  }} else {{
+    j = f.get();
+  }}
+}}
+
+inline void from_json(const nlohmann::json& j, F64& f) {{
+  if (j == "Infinity") {{
+    f.set(std::numeric_limits<double>::infinity());
+  }} else if (j == "-Infinity") {{
+    f.set(-std::numeric_limits<double>::infinity());
+  }} else if (j == "NaN") {{
+    f.set(std::numeric_limits<double>::quiet_NaN());
+  }} else {{
+    f.set(j.get<double>());
+  }}
+}}
+
+{chr(10).join(cpp_type_decls)}
+{"".join(cpp_enum_defs.values())}
+{"".join(dict(sorted(cpp_class_defs.items(), key=lambda x: class_ordering[x[0]])).values())}
+{chr(10).join(cpp_json_defs)}
+
+template <typename T> ForwardRef<T>::ForwardRef(ForwardRef<T>&&) = default;
+template <typename T> ForwardRef<T>& ForwardRef<T>::operator=(ForwardRef<T>&&) = default;
+template <typename T> ForwardRef<T>::~ForwardRef() = default;
+}} // namespace _export
+}} // namespace torch
+"""
+    thrift_schema = f"""
+namespace py3 torch._export
+namespace cpp2 torch._export.schema
+{chr(10).join(thrift_enum_defs)}
+{chr(10).join(dict(sorted(thrift_type_defs.items(), key=lambda x: class_ordering[x[0]])).values())}
+"""
+    return yaml_ret, cpp_header, thrift_schema
+
+
+def _diff_schema(dst, src):
+    additions = {key: src[key] for key in src.keys() - dst.keys()}
+    subtractions = {key: dst[key] for key in dst.keys() - src.keys()}
+
+    common_keys = src.keys() & dst.keys()
+
+    versions = {"SCHEMA_VERSION", "TREESPEC_VERSION"}
+    common_keys -= versions
+
+    for key in common_keys:
+        src_kind = src[key]["kind"]
+        src_fields = src[key]["fields"]
+        dst_kind = dst[key]["kind"]
+        dst_fields = dst[key]["fields"]
+        _check(
+            src_kind == dst_kind,
+            f"Type {key} changed kind from {dst_kind} to {src_kind}",
+        )
+        if not isinstance(src_fields, dict) or not isinstance(dst_fields, dict):
+            raise AssertionError(
+                f"expected dict fields, got src={type(src_fields)}, dst={type(dst_fields)}"
+            )
+        added_fields = {
+            key: src_fields[key] for key in src_fields.keys() - dst_fields.keys()
+        }
+        subtracted_fields = {
+            key: dst_fields[key] for key in dst_fields.keys() - src_fields.keys()
+        }
+        common_fields = src_fields.keys() & dst_fields.keys()
+
+        for field in common_fields:
+            src_field = src_fields[field]
+            dst_field = dst_fields[field]
+            if src_kind == "struct":
+                _check(
+                    src_field["type"] == dst_field["type"],
+                    f"Type of the field {key}.{field} changed from {dst_field['type']} to {src_field['type']}",
+                )
+                if "default" in src_field and "default" not in dst_field:
+                    added_fields[field] = {}
+                    added_fields[field]["default"] = src_field["default"]
+                if "default" not in src_field and "default" in dst_field:
+                    subtracted_fields[field] = {}
+                    subtracted_fields[field]["default"] = dst_field["default"]
+            elif src_kind == "enum":
+                _check(
+                    src_field == dst_field,
+                    f"Value of the enum field {key}.{field} changed from {dst_field} to {src_field}",
+                )
+            elif src_kind == "union":
+                _check(
+                    src_field["type"] == dst_field["type"],
+                    f"Type of the field {key}.{field} changed from {dst_field['type']} to {src_field['type']}",
+                )
+            else:
+                raise AssertionError(f"Unknown kind {src_kind}: {key}")
+        if len(added_fields) > 0:
+            if key in additions:
+                raise AssertionError(f"key {key} already in additions")
+            additions[key] = {}
+            additions[key]["fields"] = added_fields
+        if len(subtracted_fields) > 0:
+            if key in subtractions:
+                raise AssertionError(f"key {key} already in subtractions")
+            subtractions[key] = {}
+            subtractions[key]["fields"] = subtracted_fields
+
+    return additions, subtractions
+
+
+def _hash_content(s: str):
+    return hashlib.sha256(s.strip().encode("utf-8")).hexdigest()
+
+
+def _generate_enum_converters() -> str:
+    """Generate C++ converter functions from serialized enum values to c10 enums."""
+
+    def validate_mapping(
+        enum_class: type[IntEnum],
+        mapping: dict[int, str],
+        enum_name: str,
+        skip_values: set[int],
+    ) -> None:
+        """Validate that all enum values have corresponding c10 mappings."""
+        for member in enum_class:
+            if member.value in skip_values:
+                continue
+            if member.value not in mapping:
+                raise SchemaUpdateError(
+                    f"{enum_name}.{member.name} (value={member.value}) is missing "
+                    f"from {enum_name.upper()}_TO_C10 mapping in schema.py. "
+                    f"Please add the mapping to the c10 enum name."
+                )
+
+    # Validate that all enum values have mappings (except UNKNOWN values)
+    validate_mapping(
+        schema.ScalarType,
+        schema.SCALAR_TYPE_TO_C10,
+        "ScalarType",
+        {schema.ScalarType.UNKNOWN},
+    )
+    validate_mapping(
+        schema.Layout,
+        schema.LAYOUT_TO_C10,
+        "Layout",
+        {schema.Layout.Unknown},
+    )
+    validate_mapping(
+        schema.MemoryFormat,
+        schema.MEMORY_FORMAT_TO_C10,
+        "MemoryFormat",
+        {schema.MemoryFormat.Unknown},
+    )
+
+    def generate_converter(
+        name: str,
+        c10_type: str,
+        mapping: dict[int, str],
+        max_value: int,
+    ) -> str:
+        lines: list[str] = []
+        for i in range(max_value + 1):
+            if i in mapping:
+                lines.append(
+                    f"      static_cast<int>(c10::{c10_type}::{mapping[i]}), // {i}"
+                )
+            else:
+                lines.append(f"      kInvalid, // {i}")
+
+        return f"""
+inline c10::{c10_type} convertSerialized{name}(int serialized_value) {{
+  constexpr int kInvalid = -1;
+  constexpr int k{name}Map[] = {{
+{chr(10).join(lines)}
+  }};
+  constexpr int kMapSize = sizeof(k{name}Map) / sizeof(k{name}Map[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized {name} value out of range: ",
+      serialized_value);
+  int result = k{name}Map[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized {name} value: ",
+      serialized_value);
+  return static_cast<c10::{c10_type}>(result);
+}}
+"""
+
+    scalar_type_converter = generate_converter(
+        "ScalarType",
+        "ScalarType",
+        schema.SCALAR_TYPE_TO_C10,
+        max(schema.SCALAR_TYPE_TO_C10.keys()),
+    )
+    layout_converter = generate_converter(
+        "Layout",
+        "Layout",
+        schema.LAYOUT_TO_C10,
+        max(schema.LAYOUT_TO_C10.keys()),
+    )
+    memory_format_converter = generate_converter(
+        "MemoryFormat",
+        "MemoryFormat",
+        schema.MEMORY_FORMAT_TO_C10,
+        max(schema.MEMORY_FORMAT_TO_C10.keys()),
+    )
+
+    return f"""
+#pragma once
+
+#include <c10/core/Layout.h>
+#include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+
+// Converter functions from serialized enum values (torch._export.serde.schema)
+// to c10 enums. The serialized format has different enum values than c10.
+
+namespace torch::aot_inductor {{
+{scalar_type_converter}
+{layout_converter}
+{memory_format_converter}
+}} // namespace torch::aot_inductor
+"""
+
+
+@dataclasses.dataclass
+class _Commit:
+    result: dict[str, Any]
+    checksum_next: str
+    yaml_path: str
+    additions: dict[str, Any]
+    subtractions: dict[str, Any]
+    base: dict[str, Any]
+    checksum_head: str | None
+    cpp_header: str
+    cpp_header_path: str
+    enum_converter_header: str
+    enum_converter_header_path: str
+    thrift_checksum_head: str | None
+    thrift_checksum_real: str | None
+    thrift_checksum_next: str
+    thrift_schema: str
+    thrift_schema_path: str
+
+
+def update_schema():
+    import importlib.resources
+
+    # pyrefly: ignore [bad-argument-type]
+    if importlib.resources.is_resource(__package__, "schema.yaml"):
+        # pyrefly: ignore [bad-argument-type]
+        content = importlib.resources.read_text(__package__, "schema.yaml")
+        match = re.search("checksum<<([A-Fa-f0-9]{64})>>", content)
+        _check(match is not None, "checksum not found in schema.yaml")
+        if match is None:
+            raise AssertionError("checksum not found in schema.yaml")
+        checksum_head = match.group(1)
+
+        thrift_content = importlib.resources.read_text(
+            # pyrefly: ignore [bad-argument-type]
+            __package__,
+            "export_schema.thrift",
+        )
+        match = re.search("checksum<<([A-Fa-f0-9]{64})>>", thrift_content)
+        _check(match is not None, "checksum not found in export_schema.thrift")
+        if match is None:
+            raise AssertionError("checksum not found in export_schema.thrift")
+        thrift_checksum_head = match.group(1)
+        thrift_content = thrift_content.splitlines()
+        if not thrift_content[0].startswith("// @" + "generated"):
+            raise AssertionError(
+                f"expected first line to start with '// @generated', got {thrift_content[0]!r}"
+            )
+        if not thrift_content[1].startswith("// checksum<<"):
+            raise AssertionError(
+                f"expected second line to start with '// checksum<<', got {thrift_content[1]!r}"
+            )
+        thrift_checksum_real = _hash_content("\n".join(thrift_content[2:]))
+
+        from yaml import load, Loader
+
+        dst = load(content, Loader=Loader)
+        if not isinstance(dst, dict):
+            raise AssertionError(f"expected dict from yaml, got {type(dst)}")
+    else:
+        checksum_head = None
+        thrift_checksum_head = None
+        thrift_checksum_real = None
+        dst = {"SCHEMA_VERSION": None, "TREESPEC_VERSION": None}
+
+    src, cpp_header, thrift_schema = _staged_schema()
+    enum_converter_header = _generate_enum_converters()
+    additions, subtractions = _diff_schema(dst, src)
+    # pyrefly: ignore [missing-attribute]
+    yaml_path = __package__.replace(".", "/") + "/schema.yaml"
+    # pyrefly: ignore [missing-attribute]
+    thrift_schema_path = __package__.replace(".", "/") + "/export_schema.thrift"
+    torch_prefix = "torch/"
+    if not yaml_path.startswith(torch_prefix):
+        raise AssertionError(
+            f"yaml_path must start with {torch_prefix}, got {yaml_path}"
+        )
+    if not thrift_schema_path.startswith(torch_prefix):
+        raise AssertionError(
+            f"thrift_schema_path must start with {torch_prefix}, got {thrift_schema_path}"
+        )
+
+    return _Commit(
+        result=src,
+        checksum_next=_hash_content(repr(src)),
+        yaml_path=yaml_path,
+        additions=additions,
+        subtractions=subtractions,
+        base=dst,
+        checksum_head=checksum_head,
+        cpp_header=cpp_header,
+        cpp_header_path=torch_prefix + "csrc/utils/generated_serialization_types.h",
+        enum_converter_header=enum_converter_header,
+        enum_converter_header_path=torch_prefix
+        + "csrc/inductor/aoti_torch/generated_enum_converters.h",
+        thrift_checksum_head=thrift_checksum_head,
+        thrift_checksum_real=thrift_checksum_real,
+        thrift_checksum_next=_hash_content(thrift_schema),
+        thrift_schema=thrift_schema,
+        thrift_schema_path=thrift_schema_path,
+    )
+
+
+def check(commit: _Commit, force_unsafe: bool = False):
+    next_version = None
+    reason = ""
+    # Step 1: Detect major schema updates.
+    if len(commit.additions) > 0:
+        for k, v in commit.additions.items():
+            if k not in commit.base:
+                continue
+            kind = commit.result[k]["kind"]
+            fields = v["fields"]
+            for f, d in fields.items():
+                if kind == "struct" and "default" not in d:
+                    reason += (
+                        f"Field {k}.{f} is added to schema.py without a default value as an incompatible change "
+                        + "which requires major version bump.\n"
+                    )
+                    next_version = [commit.base["SCHEMA_VERSION"][0] + 1, 1]
+
+    if len(commit.subtractions) > 0:
+        for k, v in commit.subtractions.items():
+            if k not in commit.result:
+                continue
+            for f in v["fields"]:
+                reason = f"Field {k}.{f} is removed from schema.py as an incompatible change which requires major version bump.\n"
+            next_version = [commit.base["SCHEMA_VERSION"][0] + 1, 1]
+
+    if force_unsafe:
+        reason += "--force-unsafe is used."
+        next_version = commit.result["SCHEMA_VERSION"]
+    else:
+        # Step 2: Detect minor schema updates.
+        if next_version is None and len(commit.additions) > 0:
+            for k, v in commit.additions.items():
+                for f in v["fields"]:
+                    reason += (
+                        f"Field {k}.{f} is added to schema.py as an compatible change "
+                        + "which still requires minor version bump.\n"
+                    )
+            next_version = [
+                commit.base["SCHEMA_VERSION"][0],
+                commit.base["SCHEMA_VERSION"][1] + 1,
+            ]
+        if next_version is None and len(commit.subtractions) > 0:
+            for k, v in commit.subtractions.items():
+                for f in v["fields"]:
+                    reason += (
+                        f"Field {k}.{f} is removed from schema.py as an compatible change "
+                        + "which still requires minor version bump.\n"
+                    )
+            next_version = [
+                commit.base["SCHEMA_VERSION"][0],
+                commit.base["SCHEMA_VERSION"][1] + 1,
+            ]
+
+    return next_version, reason

kronos-venv/lib/python3.12/site-packages/torch/_export/serde/union.py

@@ -0,0 +1,104 @@
+# mypy: allow-untyped-defs
+import functools
+from collections.abc import Hashable
+from dataclasses import dataclass, fields
+from typing import TypeVar
+from typing_extensions import dataclass_transform
+
+
+T = TypeVar("T", bound="_Union")
+
+
+class _UnionTag(str):
+    __slots__ = ("_cls",)
+    _cls: Hashable
+
+    @staticmethod
+    def create(t, cls):
+        tag = _UnionTag(t)
+        if hasattr(tag, "_cls"):
+            raise AssertionError("tag already has _cls attribute")
+        tag._cls = cls
+        return tag
+
+    def __eq__(self, cmp) -> bool:
+        if not isinstance(cmp, str):
+            raise AssertionError(f"expected str, got {type(cmp)}")
+        other = str(cmp)
+        if other not in _get_field_names(self._cls):
+            raise AssertionError(
+                f"{other} is not a valid tag for {self._cls}. Available tags: {_get_field_names(self._cls)}"
+            )
+        return str(self) == other
+
+    def __hash__(self):
+        return hash(str(self))
+
+
+@functools.cache
+def _get_field_names(cls) -> set[str]:
+    return {f.name for f in fields(cls)}
+
+
+# If you turn a schema class that inherits from union into a dataclass, please use
+# this decorator to configure it. It's safe, faster and allows code sharing.
+#
+# For example, _union_dataclass customizes the __eq__ method to only check the type
+# and value property instead of default implementation of dataclass which goes
+# through every field in the dataclass.
+@dataclass_transform(eq_default=False)
+def _union_dataclass(cls: type[T]) -> type[T]:
+    if not issubclass(cls, _Union):
+        raise AssertionError(f"{cls} must inherit from {_Union}.")
+    return dataclass(repr=False, eq=False)(cls)
+
+
+class _Union:
+    _type: _UnionTag
+
+    @classmethod
+    def create(cls, **kwargs):
+        if len(kwargs) != 1:
+            raise AssertionError(f"expected exactly 1 kwarg, got {len(kwargs)}")
+        obj = cls(**{**{f.name: None for f in fields(cls)}, **kwargs})  # type: ignore[arg-type]
+        obj._type = _UnionTag.create(next(iter(kwargs.keys())), cls)
+        return obj
+
+    def __post_init__(self):
+        if any(
+            f.name in ("type", "_type", "create", "value")
+            for f in fields(self)  # type: ignore[arg-type, misc]
+        ):
+            raise AssertionError(
+                "field names 'type', '_type', 'create', 'value' are reserved"
+            )
+
+    @property
+    def type(self) -> str:
+        try:
+            return self._type
+        except AttributeError as e:
+            raise RuntimeError(
+                f"Please use {type(self).__name__}.create to instantiate the union type."
+            ) from e
+
+    @property
+    def value(self):
+        return getattr(self, self.type)
+
+    def __getattribute__(self, name):
+        attr = super().__getattribute__(name)
+        if attr is None and name in _get_field_names(type(self)) and name != self.type:  # type: ignore[arg-type]
+            raise AttributeError(f"Field {name} is not set.")
+        return attr
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _Union):
+            return False
+        return self.type == other.type and self.value == other.value
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __repr__(self):
+        return f"{type(self).__name__}({self.type}={getattr(self, self.type)})"

kronos-venv/lib/python3.12/site-packages/torch/_export/tools.py

@@ -0,0 +1,148 @@
+# mypy: allow-untyped-defs
+import logging
+import warnings
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.export
+import torch.export._trace
+from torch._utils_internal import log_export_usage
+
+
+log = logging.getLogger(__name__)
+
+__all__ = ["report_exportability"]
+
+
+def _generate_inputs_for_submodules(
+    model: torch.nn.Module,
+    target_submodules: Iterable[str],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
+) -> dict[str, tuple[Any, Any]]:
+    """
+    Generate inputs for targeting submdoules in the given model. Note that if two submodules refer to the same obj, this
+    function doesn't work.
+
+    Args:
+        model: root model.
+        inputs: inputs to the root model.
+        target_submodules: submodules that we want to generate inputs for.
+
+    Returns:
+        A dict that maps from submodule name to its inputs.
+    """
+    kwargs = kwargs or {}
+
+    handles = []
+    results = {}
+    submodule_to_names = {mod: name for name, mod in model.named_modules()}
+
+    def pre_forward(module, module_args, module_kwargs):
+        results[submodule_to_names[module]] = (module_args, module_kwargs)
+
+    try:
+        for name, mod in model.named_modules():
+            if name in target_submodules:
+                handles.append(
+                    mod.register_forward_pre_hook(pre_forward, with_kwargs=True)
+                )
+        model(*args, **kwargs)
+    except Exception as e:
+        warnings.warn(
+            f"Failed to generate submodule inputs because of the following error:\n{e}",
+            stacklevel=2,
+        )
+    finally:
+        for h in handles:
+            h.remove()
+    return results
+
+
+def report_exportability(
+    mod: torch.nn.Module,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
+    *,
+    strict: bool = True,
+    pre_dispatch: bool = False,
+) -> dict[str, Exception | None]:
+    """
+    Report exportability issues for a module in one-shot.
+
+    Args:
+        mod: root module.
+        args: args to the root module.
+        kwargs: kwargs to the root module.
+    Returns:
+        A dict that maps from submodule name to the exception that was raised when trying to export it.
+        `None` means the module is exportable without issue.
+    Sample output:
+        {
+            '': UnsupportedOperatorException(func=<OpOverload(op='testlib.op_missing_meta', overload='default')>),
+            'submod_1': UnsupportedOperatorException(func=<OpOverload(op='testlib.op_missing_meta', overload='default')>),
+            'submod_2': None
+        }
+    """
+
+    log_export_usage(event="export.report_exportability")
+
+    kwargs = kwargs or {}
+
+    all_submod_names = [name for name, _ in mod.named_modules() if name != ""]
+    submod_inputs = _generate_inputs_for_submodules(mod, all_submod_names, args, kwargs)
+
+    tried_module_types = set()
+    report: dict[str, Exception | None] = {}
+
+    def try_export(module, module_name, args, kwargs):
+        nonlocal submod_inputs, report, strict, pre_dispatch, tried_module_types
+
+        if type(module) in tried_module_types:
+            return
+        tried_module_types.add(type(module))
+
+        if args is not None or kwargs is not None:
+            try:
+                torch.export._trace._export(
+                    module,
+                    args,
+                    kwargs,
+                    strict=strict,
+                    pre_dispatch=pre_dispatch,
+                )
+                report[module_name] = None
+                log.info("Successfully exported `%s`", module_name)
+                return
+            except Exception as e:
+                short_msg = repr(e).split("\n")[0]
+                log.warning(
+                    "Failed exporting `%s` with exception: %s", module_name, short_msg
+                )
+                report[module_name] = e
+
+        for name, submod in module.named_children():
+            sub_module_name = name if module_name == "" else f"{module_name}.{name}"
+
+            submod_args, submod_kwargs = submod_inputs.get(
+                sub_module_name, (None, None)
+            )
+
+            try_export(submod, sub_module_name, submod_args, submod_kwargs)
+
+        return
+
+    try_export(mod, "", args, kwargs)
+
+    unique_issues = set()
+    for exception in report.values():
+        if exception is not None:
+            key = repr(exception).split("\\n")[0]
+            unique_issues.add(key)
+
+    log.warning("Found %d export issues:", len(unique_issues))
+    for issue in unique_issues:
+        log.warning(issue)
+
+    return report

kronos-venv/lib/python3.12/site-packages/torch/_export/verifier.py

@@ -0,0 +1,559 @@
+# mypy: allow-untyped-defs
+import inspect
+import math
+import operator
+from collections.abc import Iterable
+from typing import Any, final, TYPE_CHECKING
+
+import torch
+from torch._library.opaque_object import is_opaque_type
+from torch._ops import HigherOrderOperator, OpOverload
+from torch._subclasses.fake_tensor import FakeTensor
+from torch.export.graph_signature import (
+    CustomObjArgument,
+    InputKind,
+    SymBoolArgument,
+    SymFloatArgument,
+    SymIntArgument,
+    TensorArgument,
+    TokenArgument,
+)
+from torch.fx import GraphModule
+
+
+if TYPE_CHECKING:
+    from torch.export.exported_program import ExportedProgram
+
+
+class SpecViolationError(Exception):
+    pass
+
+
+def is_functional(op: OpOverload) -> bool:
+    return not op._schema.is_mutable
+
+
+def _check_has_fake_tensor(node: torch.fx.Node) -> None:
+    # TODO(angelayi): remove this in favor of _check_val
+    return _check_val(node)
+
+
+def _check_val(node: torch.fx.Node) -> None:
+    from torch.fx.experimental.symbolic_shapes import SymBool, SymFloat, SymInt
+
+    def _check_correct_val(val):
+        if val is None:
+            return True
+        elif isinstance(val, (int, bool, str, float)):
+            return True
+        elif isinstance(
+            val, (torch.memory_format, torch.dtype, torch.device, torch.layout)
+        ):
+            return True
+        elif isinstance(
+            val, (FakeTensor, torch.Tensor)
+        ):  # TODO(zhxchen17) Remove Tensor.
+            return True
+        elif isinstance(val, (SymInt, SymFloat, SymBool)):
+            return True
+        elif isinstance(val, CustomObjArgument):
+            return True
+        elif isinstance(val, Iterable):
+            return all(_check_correct_val(x) for x in val)
+        elif is_opaque_type(type(val)):
+            return True
+        return False
+
+    def _no_returns(op):
+        if not isinstance(op, OpOverload):
+            return False
+        return len(op._schema.returns) == 0
+
+    if "val" not in node.meta:
+        if node.op == "call_function" and _no_returns(node.target):
+            return
+        raise SpecViolationError(f"Node.meta {node.name} is missing val field.")
+
+    val = node.meta["val"]
+    if not _check_correct_val(val):
+        raise SpecViolationError(f"Node.meta {node.name} has invalid val field {val}")
+
+
+def _check_torch_fn(node: torch.fx.Node) -> None:
+    torch_fn = node.meta.get("torch_fn")
+    if torch_fn is None:
+        raise SpecViolationError(
+            f"Unable to find torch_fn metadata for node {node.name}"
+        )
+    if (
+        not isinstance(torch_fn, tuple)
+        and isinstance(torch_fn[0], str)
+        and isinstance(torch_fn[1], str)
+    ):
+        raise SpecViolationError(
+            f"Node.meta {node.name} has invalid torch_fn field {torch_fn}"
+        )
+
+
+class _VerifierMeta(type):
+    _registry: dict[str, type["Verifier"]] = {}
+
+    def __new__(metacls, name, bases, attrs):
+        if bases:
+            if "check" in attrs or "_check_graph_module" in attrs:
+                raise SyntaxError("Overriding method check is not allowed.")
+            if "dialect" not in attrs or attrs["dialect"] == "ATEN":
+                raise AssertionError(
+                    f"subclass must define dialect != 'ATEN', got {attrs.get('dialect')}"
+                )
+        else:
+            if "check" not in attrs:
+                raise AssertionError("base class must define 'check' method")
+            if "_check_graph_module" not in attrs:
+                raise AssertionError(
+                    "base class must define '_check_graph_module' method"
+                )
+            if attrs["dialect"] != "ATEN":
+                raise AssertionError(
+                    f"base class dialect must be 'ATEN', got {attrs['dialect']}"
+                )
+
+        if not isinstance(attrs["dialect"], str):
+            raise AssertionError(f"dialect must be str, got {type(attrs['dialect'])}")
+        ret = type.__new__(metacls, name, bases, attrs)
+        metacls._registry[attrs["dialect"]] = ret  # type: ignore[assignment]
+        return ret
+
+
+def getattr_recursive(obj: Any, target: str) -> Any:
+    target_atoms = target.split(".")
+    attr_itr = obj
+    for i, atom in enumerate(target_atoms):
+        if not hasattr(attr_itr, atom):
+            raise RuntimeError(
+                f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}"
+            )
+        attr_itr = getattr(attr_itr, atom)
+    return attr_itr
+
+
+class Verifier(metaclass=_VerifierMeta):
+    dialect = "ATEN"
+
+    def allowed_builtin_ops(self) -> list:
+        return [
+            operator.getitem,
+            operator.add,
+            operator.mul,
+            operator.sub,
+            operator.truediv,
+            operator.ge,
+            operator.le,
+            operator.gt,
+            operator.lt,
+            operator.eq,
+            operator.ne,
+            operator.floordiv,
+            operator.mod,
+            operator.and_,
+            operator.or_,
+            operator.not_,
+            operator.pow,
+            operator.neg,
+            operator.abs,
+            operator.lshift,
+            operator.rshift,
+            math.ceil,
+            math.floor,
+            math.trunc,
+            round,
+        ]
+
+    def allowed_op_types(self) -> tuple[type[Any], ...]:
+        return (OpOverload, HigherOrderOperator)
+
+    def allowed_getattr_types(self) -> tuple[type[Any], ...]:
+        return (torch.fx.GraphModule, torch.utils._pytree.TreeSpec)
+
+    def allowed_getattr_types_for_subgm(self) -> tuple[type[Any], ...]:
+        # subgm in HOP's argument could has have getattr(weight) nodes, thus stateful
+        return (
+            torch.fx.GraphModule,
+            torch.nn.parameter.Parameter,
+            torch.Tensor,  # for buffer and constant tensor
+            torch.utils._pytree.TreeSpec,
+        )
+
+    def check_valid_op(self, op):
+        pass
+
+    def check_additional(self, gm: GraphModule) -> None:
+        """
+        Additional checks that are specific to some dialects.
+        """
+
+    @final
+    def check(self, ep: "ExportedProgram") -> None:
+        self._check_graph_module(ep.graph_module)
+        _verify_exported_program_module_call_graph(ep)
+        _verify_exported_program_signature(ep)
+
+    @final
+    def _check_graph_module(self, gm: torch.fx.GraphModule) -> None:
+        def _allowed_getattr_types(is_toplevel_gm) -> tuple[type[Any], ...]:
+            if is_toplevel_gm:
+                ret = self.allowed_getattr_types()
+            else:
+                ret = self.allowed_getattr_types_for_subgm()
+            if any(t is object for t in ret):
+                raise AssertionError("allowed_getattr_types must not contain 'object'")
+            return ret
+
+        def _check_valid_op(op) -> None:
+            def _allowed_builtin_ops() -> list:
+                ret = self.allowed_builtin_ops()
+                if not all(inspect.isbuiltin(op) for op in ret):
+                    raise AssertionError("allowed_builtin_ops must all be builtins")
+                return ret
+
+            def _allowed_op_types() -> tuple[type[Any], ...]:
+                ret = self.allowed_op_types()
+                if any(t is object for t in ret):
+                    raise AssertionError("allowed_op_types must not contain 'object'")
+                return ret
+
+            # TODO Remove this allowlist.
+            _allowed_torch_functions = (
+                torch.autograd.grad_mode.set_grad_enabled,
+                torch.sym_int,
+                torch.sym_float,
+                torch.sym_ite,
+                torch.sym_max,
+                torch.sym_min,
+                torch.sym_not,
+                torch.sym_sqrt,
+                torch.sym_sum,
+                torch.export.custom_ops._call_custom_autograd_function_in_pre_dispatch,
+                # TODO (tmanlaibaatar)
+                # Predispatch export is able to contain autograd ops.
+                # These will be modeled as HOO later
+                torch._C._set_grad_enabled,
+                torch.amp.autocast_mode._enter_autocast,
+                torch.amp.autocast_mode._exit_autocast,
+                torch.fx.experimental.symbolic_shapes.cast_symbool_to_symint_guardless,
+                torch._functorch.predispatch._add_batch_dim,
+                torch._functorch.predispatch._remove_batch_dim,
+                torch._functorch.predispatch._vmap_increment_nesting,
+                torch._functorch.predispatch._vmap_decrement_nesting,
+                torch._functorch.predispatch.lazy_load_decompositions,
+            )
+
+            if not isinstance(op, _allowed_op_types()):
+                if (
+                    op not in _allowed_builtin_ops()
+                    and op not in _allowed_torch_functions
+                ):
+                    raise SpecViolationError(
+                        f"Operator '{op}' is not an allowed operator type: {_allowed_op_types()}\n"
+                        f"Valid builtin ops: {_allowed_builtin_ops()}"
+                        f"Valid torch functions: {_allowed_torch_functions}"
+                    )
+
+            if isinstance(op, OpOverload):
+                # All ops functional
+                # TODO (tmanlaibaatar) more proper way is needed here
+                if self.dialect != "TRAINING" and not is_functional(op):
+                    raise SpecViolationError(f"operator '{op}' is not functional")
+            self.check_valid_op(op)
+
+        for mod in gm.modules():
+            is_toplevel_gm = mod is gm
+
+            if not isinstance(mod, torch.fx.GraphModule):
+                continue
+
+            mod.graph.lint()
+            for node in mod.graph.nodes:
+                # TODO(T140410192): should have fake tensor for all dialects
+                if node.op in {"call_module", "call_method"}:
+                    raise SpecViolationError(
+                        f"call_module is not valid: got a class '{node.target}' ",
+                    )
+
+                elif node.op == "call_function":
+                    _check_val(node)
+
+                    _check_valid_op(node.target)
+
+                elif node.op == "get_attr":
+                    if not isinstance(node.target, str):
+                        raise SpecViolationError(
+                            f"Expected get_attr target to be string, but got {type(node.target)}"
+                        )
+
+                    attr = getattr_recursive(mod, node.target)
+                    if isinstance(attr, torch.nn.Module):
+
+                        def _is_type(name, ty):
+                            return isinstance(getattr(attr, name, None), ty)
+
+                        if type(attr).__name__ == "LoweredBackendModule":
+                            if (
+                                _is_type("backend_id", str)
+                                and hasattr(attr, "original_module")
+                                and hasattr(attr, "module_name")
+                                and getattr(attr, "backend_id", None) == "aoti"
+                            ):
+                                continue
+                            if (
+                                _is_type("backend_id", str)
+                                and _is_type("processed_bytes", bytes)
+                                and _is_type("compile_specs", list)
+                                and hasattr(attr, "original_module")
+                            ):
+                                continue
+                            else:
+                                backend_id = getattr(attr, "backend_id", None)
+                                processed_bytes = getattr(attr, "processed_bytes", None)
+                                compile_specs = getattr(attr, "compile_specs", None)
+                                raise SpecViolationError(
+                                    f"Invalid get_attr type {type(attr)}. \n"
+                                    f"LoweredBackendModule fields: "
+                                    f"backend_id(str) : {type(backend_id)}, "
+                                    f"processed_bytes(bytes) : {type(processed_bytes)}, "
+                                    f"compile_specs(list) : {type(compile_specs)}"
+                                )
+                        elif type(attr).__name__ == "AOTInductorEPModule":
+                            continue
+
+                        elif type(attr).__name__ == "AOTInductorRunnerWrapper":
+                            continue
+
+                    if not isinstance(attr, _allowed_getattr_types(is_toplevel_gm)):
+                        raise SpecViolationError(
+                            f"Invalid get_attr type {type(attr)} on target {node.target}. \n"
+                            f"Valid get_attr types: {_allowed_getattr_types(is_toplevel_gm)}"
+                        )
+
+                elif node.op == "placeholder":
+                    _check_val(node)
+                # TODO(zhxchen17)
+                # elif node.op == "output":
+                #     _check_flattened_outputs()
+
+        self.check_additional(gm)
+
+
+class TrainingIRVerifier(Verifier):
+    dialect = "TRAINING"
+
+
+def _verify_exported_program_module_call_graph(exported_program) -> None:
+    module_call_graph = exported_program.module_call_graph
+    nodes = {node.name for node in exported_program.graph.nodes}
+    for entry in module_call_graph:
+        if entry.signature is not None:
+            for arg in entry.signature.inputs:
+                if arg.name and arg.name not in nodes:
+                    raise SpecViolationError(
+                        f"Input {arg.name} does not exist in the graph."
+                    )
+            for arg in entry.signature.outputs:
+                if arg.name and arg.name not in nodes:
+                    raise SpecViolationError(
+                        f"Output {arg.name} does not exist in the graph."
+                    )
+
+
+def _verify_exported_program_signature(exported_program) -> None:
+    # Check ExportedProgram signature matches
+    gs = exported_program.graph_signature
+
+    # Check every node in the signature exists in the graph
+    input_node_names = [
+        node.name for node in exported_program.graph.nodes if node.op == "placeholder"
+    ]
+
+    if len(input_node_names) != len(gs.input_specs):
+        input_spec_names = [
+            spec.arg.name for spec in gs.input_specs if hasattr(spec.arg, "name")
+        ]
+        missing_in_specs = set(input_node_names) - set(input_spec_names)
+        missing_in_graph = set(input_spec_names) - set(input_node_names)
+        raise SpecViolationError(
+            f"Number of graph inputs ({len(input_node_names)}) "
+            f"does not match number of inputs in the graph signature ({len(gs.input_specs)})\n"
+            f"Placeholders missing input_specs: {missing_in_specs}\n"
+            f"Input_specs missing placeholders: {missing_in_graph}"
+        )
+
+    for input_spec, node in zip(gs.input_specs, input_node_names):
+        if isinstance(
+            input_spec.arg,
+            (TensorArgument, SymIntArgument, SymFloatArgument, SymBoolArgument),
+        ):
+            if input_spec.arg.name != node:
+                raise SpecViolationError(
+                    f"Input spec name {input_spec.arg.name} does not match node name {node}"
+                )
+
+        if input_spec.kind == InputKind.USER_INPUT:
+            continue
+
+        elif input_spec.kind == InputKind.PARAMETER:
+            if not isinstance(input_spec.arg, TensorArgument):
+                raise SpecViolationError(
+                    f"Parameter {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
+                )
+            if input_spec.target is None:
+                raise SpecViolationError(
+                    f"InputSpec for {input_spec.name} has no target."
+                )
+
+            param = input_spec.target
+            if param not in exported_program.state_dict:
+                raise SpecViolationError(f"Parameter {param} is not in the state dict.")
+
+            if not isinstance(exported_program.state_dict[param], torch.nn.Parameter):
+                raise SpecViolationError(
+                    f"State dict entry for parameter {param} is not an instance of torch.nn.Parameter."
+                )
+
+        elif input_spec.kind == InputKind.BUFFER:
+            if not isinstance(input_spec.arg, TensorArgument):
+                raise SpecViolationError(
+                    f"Buffer {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
+                )
+            if input_spec.target is None:
+                raise SpecViolationError(
+                    f"InputSpec for {input_spec.name} has no target."
+                )
+
+            buffer = input_spec.target
+            if input_spec.persistent is None:
+                raise SpecViolationError(
+                    f"Buffer {buffer} is missing a persistence flag"
+                )
+
+            if (
+                input_spec.persistent is True
+                and buffer not in exported_program.state_dict
+            ):
+                raise SpecViolationError(f"Buffer {buffer} is not in the state dict.")
+
+            if input_spec.persistent is False and buffer in exported_program.state_dict:
+                raise SpecViolationError(
+                    f"Non-persistent buffer {buffer} is in the state dict, it should not be."
+                )
+        elif input_spec.kind == InputKind.CONSTANT_TENSOR:
+            if not isinstance(input_spec.arg, TensorArgument):
+                raise SpecViolationError(
+                    f"Constant tensor {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
+                )
+            if input_spec.target is None:
+                raise SpecViolationError(
+                    f"InputSpec for {input_spec.name} has no target."
+                )
+
+            tensor_const = input_spec.target
+            if tensor_const not in exported_program.constants:
+                raise SpecViolationError(
+                    f"Constant tensor {tensor_const} is not in the constants dictionary."
+                )
+        elif input_spec.kind == InputKind.CUSTOM_OBJ:
+            if not isinstance(input_spec.arg, CustomObjArgument):
+                raise SpecViolationError(
+                    f"Custom object {input_spec.name} is not a custom object argument. Found {input_spec.arg} instead."
+                )
+            if input_spec.target is None:
+                raise SpecViolationError(
+                    f"InputSpec for {input_spec.name} has no target."
+                )
+
+            custom_obj = input_spec.target
+            if custom_obj not in exported_program.constants:
+                raise SpecViolationError(
+                    f"Custom object {custom_obj} is not in the constants dictionary."
+                )
+        elif input_spec.kind == InputKind.TOKEN:
+            if not isinstance(input_spec.arg, TokenArgument):
+                raise SpecViolationError(
+                    f"Constant tensor {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
+                )
+        else:
+            raise SpecViolationError(f"Unknown InputKind {input_spec.kind}.")
+
+    # Check outputs
+    output_node = list(exported_program.graph.nodes)[-1]
+    if output_node.op != "output":
+        raise AssertionError(f"last node must be output, got {output_node.op}")
+    output_nodes = [
+        arg.name if isinstance(arg, torch.fx.Node) else arg
+        for arg in output_node.args[0]
+    ]
+
+    if len(output_nodes) != len(gs.output_specs):
+        output_spec_names = [
+            spec.arg.name if hasattr(spec.arg, "name") else str(spec.arg)
+            for spec in gs.output_specs
+        ]
+        missing_out_specs = set(output_nodes) - set(output_spec_names)
+        missing_out_graph = set(output_spec_names) - set(output_nodes)
+        raise SpecViolationError(
+            f"Number of output nodes {len(output_nodes)} is different "
+            f"Than the number of outputs specified by the graph signature: {len(gs.output_specs)}\n"
+            f"Nodes missing output_specs: {missing_out_specs}\n"
+            f"Output_specs missing nodes: {missing_out_graph}"
+        )
+
+    num_tokens = len(gs.output_tokens)
+    end = (
+        len(gs.buffers_to_mutate)
+        + len(gs.parameters_to_mutate)
+        + len(gs.user_inputs_to_mutate)
+        + num_tokens
+    )
+    mutate_nodes: list[str] = output_nodes[num_tokens:end]
+    user_output_nodes = output_nodes[end : end + len(gs.user_outputs)]
+
+    for mutation_node in mutate_nodes:
+        if mutation_node in gs.buffers_to_mutate:
+            if gs.buffers_to_mutate[mutation_node] not in gs.buffers:
+                raise SpecViolationError(
+                    f"Buffer output {mutation_node} does not point to a buffer that exists. \n"
+                    f"Dict of buffers that are mutated, in order: {gs.buffers_to_mutate} \n"
+                    f"Buffer nodes available: {gs.buffers} \n"
+                )
+        elif mutation_node in gs.parameters_to_mutate:
+            if gs.parameters_to_mutate[mutation_node] not in gs.parameters:
+                raise SpecViolationError(
+                    f"Parameter output {mutation_node} does not point to a parameter that exists. \n"
+                    f"Dict of parameters that are mutated, in order: {gs.parameters_to_mutate} \n"
+                    f"Parameter nodes available: {gs.parameters} \n"
+                )
+        elif mutation_node in gs.user_inputs_to_mutate:
+            if gs.user_inputs_to_mutate[mutation_node] not in gs.user_inputs:
+                raise SpecViolationError(
+                    f"User input output {mutation_node} does not point to a user input that exists. \n"
+                    f"Dict of user inputs that are mutated, in order: {gs.user_inputs_to_mutate} \n"
+                    f"User input nodes available: {gs.user_inputs} \n"
+                )
+        else:
+            raise SpecViolationError(
+                f"Mutation node {mutation_node} is neither a buffer nor a user input. "
+                f"Buffers to mutate: {gs.buffers_to_mutate}, User inputs to mutate: {gs.user_inputs_to_mutate}"
+            )
+
+    for user_output_node, user_output_name in zip(user_output_nodes, gs.user_outputs):
+        if user_output_node != user_output_name:
+            raise SpecViolationError(
+                f"User output {user_output_node} is not in the correct "
+                "order or is not found in the "
+                f"exported program's user_output list: {gs.user_outputs}. "
+            )
+
+
+def load_verifier(dialect: str) -> type[Verifier]:
+    if dialect == "ATEN" or dialect == "":
+        return _VerifierMeta._registry.get(dialect, Verifier)
+    return _VerifierMeta._registry[dialect]

kronos-venv/lib/python3.12/site-packages/torch/_export/wrappers.py

@@ -0,0 +1,354 @@
+# mypy: allow-untyped-defs
+import inspect
+from contextlib import contextmanager
+from functools import wraps
+
+import torch
+import torch._custom_ops
+from torch._C import DispatchKey
+from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
+from torch._higher_order_ops.flat_apply import (
+    _ConstantFunction,
+    flat_apply,
+    to_graphable,
+)
+from torch._higher_order_ops.strict_mode import strict_mode
+from torch._higher_order_ops.utils import autograd_not_implemented
+from torch._ops import HigherOrderOperator
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.proxy_tensor import (
+    PreDispatchTorchFunctionMode,
+    ProxyTorchDispatchMode,
+    track_tensor_tree,
+)
+from torch.utils import _pytree as pytree
+from torch.utils._python_dispatch import is_traceable_wrapper_subclass_type
+
+
+class ExportTracepoint(HigherOrderOperator):
+    def __init__(self):
+        super().__init__("_export_tracepoint")
+
+    def __call__(self, *args, **kwargs):
+        # pyrefly: ignore [missing-attribute]
+        return super().__call__(*args, **kwargs)
+
+
+_export_tracepoint = ExportTracepoint()
+
+
+@_export_tracepoint.py_impl(ProxyTorchDispatchMode)
+def export_tracepoint_dispatch_mode(mode, *args, **kwargs):
+    p_args, p_kwargs = pytree.tree_map(mode.tracer.unwrap_proxy, (args, kwargs))
+    proxy = mode.tracer.create_proxy(
+        "call_function", _export_tracepoint, p_args, p_kwargs
+    )
+    return track_tensor_tree(args, proxy, constant=None, tracer=mode.tracer)
+
+
+@_export_tracepoint.py_impl(FakeTensorMode)
+def export_tracepoint_fake_tensor_mode(mode, *args, **kwargs):
+    with mode:
+        return args
+
+
+@_export_tracepoint.py_functionalize_impl
+def export_tracepoint_functional(ctx, *args, **kwargs):
+    unwrapped_args = ctx.unwrap_tensors(args)
+    unwrapped_kwargs = ctx.unwrap_tensors(kwargs)
+
+    with ctx.redispatch_to_next():
+        _export_tracepoint(*unwrapped_args, **unwrapped_kwargs)
+        return args
+
+
+_export_tracepoint.py_impl(DispatchKey.Autograd)(
+    autograd_not_implemented(_export_tracepoint, deferred_error=True)
+)
+
+
+@_export_tracepoint.py_impl(DispatchKey.CPU)
+def export_tracepoint_cpu(*args, **kwargs):
+    return args
+
+
+def _wrap_submodule(mod, path, module_call_specs):
+    if not isinstance(mod, torch.nn.Module):
+        raise AssertionError(f"expected torch.nn.Module, got {type(mod)}")
+    if path == "":
+        raise AssertionError("path must not be empty")
+    submodule = torch.fx.graph_module._get_attr(mod, path)
+
+    def update_module_call_signatures(path, in_spec, out_spec):
+        if path in module_call_specs:
+            if module_call_specs[path]["in_spec"] != in_spec:
+                raise AssertionError(
+                    f"in_spec mismatch for {path}: {module_call_specs[path]['in_spec']} != {in_spec}"
+                )
+            if module_call_specs[path]["out_spec"] != out_spec:
+                raise AssertionError(
+                    f"out_spec mismatch for {path}: {module_call_specs[path]['out_spec']} != {out_spec}"
+                )
+        module_call_specs[path] = {"in_spec": in_spec, "out_spec": out_spec}
+
+    def check_flattened(flat_args):
+        for a in flat_args:
+            if not (isinstance(a, (torch.Tensor, str, int, float, bool)) or a is None):
+                raise AssertionError(
+                    f"Only Tensors or scalars are supported as pytree flattened inputs, got: {a}"
+                )
+
+    def pre_hook(module, args, kwargs):
+        flat_args, in_spec = pytree.tree_flatten((args, kwargs))
+        check_flattened(flat_args)
+        flat_args = _export_tracepoint(*flat_args, kind="module_call_inputs", path=path)
+        args, kwargs = pytree.tree_unflatten(flat_args, in_spec)
+        return args, kwargs
+
+    def post_hook(module, args, kwargs, res):
+        _, in_spec = pytree.tree_flatten((args, kwargs))
+        flat_res, out_spec = pytree.tree_flatten(res)
+        check_flattened(flat_res)
+        flat_res = _export_tracepoint(*flat_res, kind="module_call_outputs", path=path)
+        update_module_call_signatures(path, in_spec, out_spec)
+        return pytree.tree_unflatten(flat_res, out_spec)
+
+    pre_handle = submodule.register_forward_pre_hook(pre_hook, with_kwargs=True)
+    post_handle = submodule.register_forward_hook(post_hook, with_kwargs=True)
+    return pre_handle, post_handle
+
+
+@contextmanager
+def _wrap_submodules(f, preserve_signature, module_call_signatures):
+    handles = []
+
+    try:
+        for path in preserve_signature:
+            handles.extend(_wrap_submodule(f, path, module_call_signatures))
+        yield
+    finally:
+        for handle in handles:
+            handle.remove()
+
+
+def _mark_strict_experimental(cls):
+    def call(self, *args):
+        return strict_mode(self, args)
+
+    cls.__call__ = call
+    return cls
+
+
+def _register_func_spec_proxy_in_tracer(tracer, name, spec):
+    """
+    This is a wrapper utility method on top of tracer to cache the
+    already registered subclass spec attribute. This is useful because
+    Subclass.__init__ will be same for each subclass. By default, fx will
+    create multiple attributes/proxies for given attribute.
+    """
+    fx_name = name + "0"
+    if hasattr(tracer.root, fx_name):
+        if getattr(tracer.root, fx_name) != spec:
+            raise AssertionError(f"spec mismatch for {fx_name}")
+        return tracer.create_proxy("get_attr", fx_name, (), {})
+
+    qualname = tracer.get_fresh_qualname(name)
+    setattr(tracer.root, qualname, spec)
+    return tracer.create_proxy("get_attr", qualname, (), {})
+
+
+def _emit_flat_apply_call(
+    *,
+    tracer,
+    spec_name: str,
+    const_target_for_apply,
+    graphable_args,
+    track_value,
+    call_spec_cache_key: str,
+):
+    # Flatten to graphable form and record the spec on the FX root
+    flat_args, in_spec = to_graphable(graphable_args)
+    qualname = tracer.get_fresh_qualname(spec_name)  # type: ignore[union-attr]
+    setattr(tracer.root, qualname, in_spec)  # type: ignore[union-attr]
+    spec_proxy = tracer.create_proxy("get_attr", qualname, (), {})
+
+    # Reuse/cached ConstantFunction spec on the root
+    _, func_spec = pytree.tree_flatten(_ConstantFunction(const_target_for_apply))
+    func_spec_proxy = _register_func_spec_proxy_in_tracer(
+        tracer, f"{call_spec_cache_key}_const_func_spec", func_spec
+    )
+
+    # Map runtime args -> proxies (always via tracer.unwrap_proxy now)
+    flat_proxy_args = pytree.tree_map(tracer.unwrap_proxy, flat_args)
+
+    # Emit flat_apply and track result structure
+    out_proxy = tracer.create_proxy(
+        "call_function", flat_apply, (func_spec_proxy, spec_proxy, *flat_proxy_args), {}
+    )
+    track_tensor_tree(track_value, out_proxy, constant=None, tracer=tracer)
+
+
+def _is_init(fn):
+    return callable(fn) and fn.__name__ == "__init__"
+
+
+def mark_subclass_constructor_exportable_experimental(constructor_subclass):
+    """
+    Experimental decorator that makes subclass to be traceable in export
+    with pre-dispatch IR. To make your subclass traceble in export, you need to:
+        1. Implement __init__ method for your subclass (Look at DTensor implementation)
+        2. Decorate your __init__ method with _mark_constructor_exportable_experimental
+        3. Put torch._dynamo_disable decorator to prevent dynamo from peeking into its' impl
+
+    Example:
+
+    class FooTensor(torch.Tensor):
+        @staticmethod
+        def __new__(cls, elem, *, requires_grad=False):
+            # ...
+            return torch.Tensor._make_subclass(cls, elem, requires_grad=requires_grad)
+
+        @torch._dynamo_disable
+        @mark_subclass_constructor_exportable_experimental
+        def __init__(self, elem, ...):
+            # ...
+    """
+    if not _is_init(constructor_subclass):
+        raise RuntimeError(
+            f"torch._export.wrappers.mark_constructor_exportable_experimental can only be applied on subclass tensor.__init__"
+            f"But, you are adding it on {constructor_subclass.__name__} which is not supported. "
+            f"If __init__ doesn't exist on your subclass, please add it. Look at DTensor.__init__ implementation for example"
+        )
+
+    def wrapper(*args, **kwargs):
+        constructor_subclass(*args, **kwargs)
+
+        if not torch.compiler.is_exporting():
+            return
+
+        if not is_traceable_wrapper_subclass_type(type(args[0])):
+            if not constructor_subclass.__qualname__.endswith("__init__"):
+                raise AssertionError(
+                    f"expected __qualname__ to end with '__init__', got {constructor_subclass.__qualname__}"
+                )
+            obj_name = constructor_subclass.__qualname__[: -len("__init__")]
+            raise RuntimeError(
+                f"Can't intercept {obj_name} in export because this object is not a traceable "
+                f"tensor subclass. Please look at DTensor.__init__ implementation as an example of proper usage of this API."
+            )
+
+        mode = _maybe_find_pre_dispatch_tf_mode_for_export()
+        if mode is None:
+            return
+
+        if not isinstance(mode, PreDispatchTorchFunctionMode):
+            raise AssertionError(
+                f"expected PreDispatchTorchFunctionMode, got {type(mode)}"
+            )
+
+        tracer = mode.tracer
+        subclass = args[0]
+        graphable = (tuple(args[1:]), kwargs)
+
+        spec_name = "_".join(constructor_subclass.__qualname__.lower().split("."))
+        call_spec_cache_key = type(subclass).__name__.lower()
+
+        _emit_flat_apply_call(
+            tracer=tracer,
+            spec_name=spec_name,
+            const_target_for_apply=type(subclass),
+            graphable_args=graphable,
+            track_value=subclass,  # track the constructed subclass instance
+            call_spec_cache_key=call_spec_cache_key,
+        )
+        return
+
+    return wrapper
+
+
+def allow_in_pre_dispatch_graph(func):
+    """
+    Experimental decorator that adds user function to export pre-dispatch graph. Note that
+    we only support custom autograd function/subclass constructors today. To use this function:
+        1. For subclasses:
+            1. refer to instructions in mark_subclass_constructor_exportable_experimental
+        2. Define apply method on your custom autograd function and apply this decorator.
+
+    Example:
+
+    class MyCoolCustomAutogradFunc(autograd.Function):
+        @classmethod
+        @torch._export.wrappers.allow_in_pre_dispatch_graph
+        def apply(cls, *args, **kwargs):
+            return super(MyCoolCustomAutogradFunc, cls).apply(*args, **kwargs)
+
+    """
+    if _is_init(func):
+        return mark_subclass_constructor_exportable_experimental(func)
+
+    if not (_is_init(func) or func.__name__ == "apply"):
+        raise RuntimeError(
+            f"torch._export.wrappers.allow_in_pre_dispatch_graph can only be applied on subclass tensor.__init_ "
+            f"or custom_autograd_function.apply. "
+            f"But, you are adding it on {func.__name__} which is not supported. "
+            f"If __init__ doesn't exist on your subclass, please add it. Look at DTensor.__init__ implementation for example. "
+            f"If you are adding it on custom autograd function, please add it on apply method. "
+            f"If anything else, file an issue on github and we may consider extending our support. "
+        )
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if not torch.compiler.is_exporting():
+            return func(*args, **kwargs)
+
+        if not inspect.isclass(args[0]):
+            return func(*args, **kwargs)
+
+        if not issubclass(args[0], torch.autograd.Function):
+            return func(*args, **kwargs)
+
+        from torch._ops import _get_dispatch_mode_pre_dispatch
+
+        mode = _get_dispatch_mode_pre_dispatch(torch._C._TorchDispatchModeKey.PROXY)
+        if mode is None:
+            return func(*args, **kwargs)
+
+        # Sometimes custom autograd functions can call into HOPs that don't have proxy impl
+        # at PreDispatch level, so we just dispatch it below to get the concrete result.
+        include_to_set = torch._C._dispatch_tls_local_include_set().remove(
+            torch._C.DispatchKey.PreDispatch
+        )
+        exclude_to_set = (
+            torch._C._dispatch_tls_local_exclude_set()
+            | torch._C.DispatchKeySet(torch._C.DispatchKey.PreDispatch)
+        )
+
+        with torch._C._ForceDispatchKeyGuard(include_to_set, exclude_to_set):
+            out = func(*args, **kwargs)
+
+        if not mode.pre_dispatch:
+            raise AssertionError("Should only do this in predispatch")
+        tracer = mode.tracer
+
+        function_cls_name = f"{args[0].__module__}.{args[0].__qualname__}"
+        graphable = ((function_cls_name, *args[1:]), kwargs)
+
+        from torch.export.custom_ops import (
+            _call_custom_autograd_function_in_pre_dispatch,
+        )
+
+        spec_name = "_".join(function_cls_name.split("."))
+        call_spec_cache_key = type(
+            _call_custom_autograd_function_in_pre_dispatch
+        ).__name__.lower()
+        _emit_flat_apply_call(
+            tracer=tracer,
+            spec_name=spec_name,
+            const_target_for_apply=_call_custom_autograd_function_in_pre_dispatch,
+            graphable_args=graphable,
+            track_value=out,
+            call_spec_cache_key=call_spec_cache_key,
+        )
+        return out
+
+    return wrapper