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

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/_lazy_import.py

@@ -0,0 +1,39 @@
+"""Utility to lazily import modules."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any, TYPE_CHECKING
+
+
+class _LazyModule:
+    """Lazily import a module."""
+
+    def __init__(self, module_name: str) -> None:
+        self._name = module_name
+        self._module: Any = None
+
+    def __repr__(self) -> str:
+        return f"<lazy module '{self._name}'>"
+
+    def __getattr__(self, attr: str) -> object:
+        if self._module is None:
+            self._module = importlib.import_module(".", self._name)
+        return getattr(self._module, attr)
+
+
+# Import the following modules during type checking to enable code intelligence features,
+# such as auto-completion in tools like pylance, even when these modules are not explicitly
+# imported in user code.
+# NOTE: Add additional used imports here.
+if TYPE_CHECKING:
+    import onnx
+    import onnx_ir  # type: ignore[import-untyped, import-not-found]
+    import onnxscript
+    import onnxscript._framework_apis.torch_2_11 as onnxscript_apis
+
+else:
+    onnx = _LazyModule("onnx")
+    onnx_ir = _LazyModule("onnx_ir")
+    onnxscript = _LazyModule("onnxscript")
+    onnxscript_apis = _LazyModule("onnxscript._framework_apis.torch_2_11")

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_analysis.py

@@ -0,0 +1,271 @@
+"""Compatibility analyzer for PyTorch models."""
+
+# mypy: allow-untyped-defs
+# flake8: noqa: B950 We do not need flake8 as it complains line length
+from __future__ import annotations
+
+import dataclasses
+import operator
+import textwrap
+import traceback
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import torch
+import torch._export.serde.schema
+from torch.export import graph_signature
+from torch.onnx._internal.exporter import _dispatching, _registration
+
+
+if TYPE_CHECKING:
+    import torch.fx
+
+
+@dataclasses.dataclass
+class ModelInfo:
+    """Information about the model."""
+
+    parameter_count: defaultdict[torch.dtype, int] = dataclasses.field(
+        default_factory=lambda: defaultdict(int)
+    )
+    buffer_count: defaultdict[torch.dtype, int] = dataclasses.field(
+        default_factory=lambda: defaultdict(int)
+    )
+    fx_node_count: int = 0
+    fx_node_op_count: defaultdict[str, int] = dataclasses.field(
+        default_factory=lambda: defaultdict(int)
+    )
+    fx_node_target_count: defaultdict[str, int] = dataclasses.field(
+        default_factory=lambda: defaultdict(int)
+    )
+    dispatch_failures: list[tuple[torch.fx.Node, str]] = dataclasses.field(
+        default_factory=list
+    )
+    inputs: dict[str, torch._export.serde.schema.TensorMeta] = dataclasses.field(
+        default_factory=dict
+    )
+    outputs: dict[str, torch._export.serde.schema.TensorMeta] = dataclasses.field(
+        default_factory=dict
+    )
+
+
+def _count_weights(
+    exported_program: torch.export.ExportedProgram,
+) -> tuple[defaultdict[torch.dtype, int], defaultdict[torch.dtype, int]]:
+    """Count the size of the parameters in the exported program."""
+
+    parameter_count: defaultdict[torch.dtype, int] = defaultdict(int)
+    buffer_count: defaultdict[torch.dtype, int] = defaultdict(int)
+    for parameter in exported_program.parameters():
+        dtype = parameter.dtype
+        parameter_count[dtype] += parameter.numel()
+
+    for buffer in exported_program.buffers():
+        dtype = buffer.dtype
+        buffer_count[dtype] += buffer.numel()
+
+    return parameter_count, buffer_count
+
+
+def _format_model_info(model_info: ModelInfo) -> str:
+    """Format the information about the model."""
+    lines = [
+        textwrap.dedent(
+            f"""\
+            PyTorch ONNX Conversion Analysis
+
+            ## Model Information
+
+            The model has {sum(model_info.parameter_count.values())} parameters and {sum(model_info.buffer_count.values())} buffers (non-trainable parameters).
+            Number of parameters per dtype:
+            ```python
+            {model_info.parameter_count}
+            ```
+            Number of buffers per dtype:
+            ```python
+            {model_info.buffer_count}
+            ```
+            """
+        ),
+        "Inputs:",
+        *[f"- `{name}`: `{meta}`" for name, meta in model_info.inputs.items()],
+        "",
+        "Outputs:",
+        *[f"- `{name}`: `{meta}`" for name, meta in model_info.outputs.items()],
+        "",
+        f"The FX graph has {model_info.fx_node_count} nodes in total. Number of FX nodes per op:",
+    ]
+    for op, count in model_info.fx_node_op_count.items():
+        lines.append(f"- `{op}`: {count}")
+    lines.append("\n")
+    lines.append("Of the call_function nodes, the counts of operators used are:\n")
+    sorted_targets = sorted(
+        model_info.fx_node_target_count.items(),
+        key=operator.itemgetter(1),
+        reverse=True,
+    )
+    for target, count in sorted_targets:
+        lines.append(f"- `{target}`: {count}")
+
+    lines.append("")
+    lines.append("## ONNX Conversion Information")
+    lines.append("")
+
+    if model_info.dispatch_failures:
+        lines.append(
+            "The model contains operators the dispatcher could not find registered ONNX decompositions for. "
+            "This may be due to missing implementations, decompositions not registered "
+            "correctly, or a bug in the dispatcher."
+        )
+        lines.append("")
+        lines.append("Errors grouped by operator:\n")
+
+        target_to_nodes = defaultdict(list)
+        for node, _ in model_info.dispatch_failures:
+            target_to_nodes[str(node.target)].append(node)
+
+        target_to_messages = {}
+        for node, message in model_info.dispatch_failures:
+            if str(node.target) not in target_to_messages:
+                target_to_messages[str(node.target)] = message
+
+        for target, nodes in sorted(
+            target_to_nodes.items(), key=operator.itemgetter(0), reverse=True
+        ):
+            message = textwrap.indent(
+                f"{target_to_messages[target]}. Example node: `{nodes[0].format_node()}`. All nodes: `{nodes}`",
+                "    ",
+            )
+            lines.append(f"- `{target}`: {message}")
+    else:
+        lines.append("All operators in the model have registered ONNX decompositions.")
+
+    return "\n".join(lines)
+
+
+def _get_io_specs(exported_program: torch.export.ExportedProgram) -> tuple[dict, dict]:
+    """Get the input and output specs of the exported program."""
+
+    nodes: dict[str, torch.fx.Node] = {
+        node.name: node for node in exported_program.graph.nodes
+    }
+    user_inputs = [
+        spec
+        for spec in exported_program.graph_signature.input_specs
+        if spec.kind == graph_signature.InputKind.USER_INPUT
+    ]
+    user_outputs = [
+        spec
+        for spec in exported_program.graph_signature.output_specs
+        if spec.kind == graph_signature.OutputKind.USER_OUTPUT
+    ]
+    inputs: dict[str, torch._export.serde.schema.TensorMeta | str] = {}
+    outputs: dict[str, torch._export.serde.schema.TensorMeta | str] = {}
+    for spec in user_inputs:
+        inputs = _log_spec_into_io_specs(spec, nodes, inputs)
+    for spec in user_outputs:
+        outputs = _log_spec_into_io_specs(spec, nodes, outputs)
+    return inputs, outputs
+
+
+def _log_spec_into_io_specs(
+    spec: graph_signature.InputSpec,
+    nodes: dict[str, torch.fx.Node],
+    inputs_or_outputs: dict[str, torch._export.serde.schema.TensorMeta | str],
+) -> dict[str, torch._export.serde.schema.TensorMeta | str]:
+    # If dynamic is set to a constant input, it becomes a
+    # symbolic argument, which is not a tensor.
+    if isinstance(spec.arg, graph_signature.ConstantArgument):
+        # Constant input does not have tensor_meta.
+        return inputs_or_outputs
+    # Symbolic arguments are not tensors, so it does not have tensor_meta,
+    # but we need to provide a string representation for them to inform users.
+    name = spec.arg.name
+    if isinstance(
+        spec.arg,
+        (
+            graph_signature.SymIntArgument,
+            graph_signature.SymFloatArgument,
+            graph_signature.SymBoolArgument,
+        ),
+    ):
+        argument_to_str: dict[type[graph_signature.ArgumentSpec], str] = {
+            graph_signature.SymIntArgument: "SymInt",
+            graph_signature.SymFloatArgument: "SymFloat",
+            graph_signature.SymBoolArgument: "SymBool",
+        }
+        inputs_or_outputs[name] = argument_to_str[type(spec.arg)]
+        return inputs_or_outputs
+    # FIXME: tensor_meta is None sometimes when the exported program still knows the shape/type
+    inputs_or_outputs[name] = nodes[name].meta["tensor_meta"]
+    return inputs_or_outputs
+
+
+def _count_fx_targets(
+    exported_program: torch.export.ExportedProgram,
+) -> defaultdict[str, int]:
+    """Count the number of targets for each node in the exported program."""
+    fx_node_target_count: defaultdict[str, int] = defaultdict(int)
+    for node in exported_program.graph.nodes:
+        if node.op == "call_function":
+            fx_node_target_count[str(node.target)] += 1
+    return fx_node_target_count
+
+
+def analyze(
+    exported_program: torch.export.ExportedProgram,
+    registry: _registration.ONNXRegistry | None = None,
+    file=None,
+) -> None:
+    """Analyze the compatibility of the exported program."""
+    # Get basic information about the model
+    model_info = ModelInfo()
+    model_info.parameter_count, model_info.buffer_count = _count_weights(
+        exported_program
+    )
+    model_info.fx_node_count = len(exported_program.graph.nodes)
+    model_info.fx_node_target_count = _count_fx_targets(exported_program)
+    inputs, outputs = _get_io_specs(exported_program)
+    model_info.inputs = inputs
+    model_info.outputs = outputs
+
+    if registry is None:
+        registry = _registration.ONNXRegistry.from_torchlib()
+
+    # Try to find ops for every node in the graph
+    for node in exported_program.graph.nodes:
+        model_info.fx_node_op_count[node.op] += 1
+        if node.op == "call_function":
+            try:
+                onnx_function, message = _dispatching.dispatch(node, registry)
+            except Exception as e:
+                message = "Critical Error in dispatcher:\n"
+                formatted_exception = "\n".join(
+                    traceback.format_exception(type(e), e, e.__traceback__)
+                )
+                message += f"```pytb\n{formatted_exception}\n```\n"
+                onnx_function = None
+            if onnx_function is None:
+                model_info.dispatch_failures.append((node, message))
+
+    # Print the results
+    report = _format_model_info(model_info)
+    print(report, file=file, flush=True)
+
+
+def compare_ops(
+    program_a: torch.export.ExportedProgram, program_b: torch.export.ExportedProgram
+) -> tuple[set[str], set[str]]:
+    """Compare and get unique ops in two exported programs.
+
+    Args:
+        program_a: The first exported program.
+        program_b: The second exported program.
+
+    Returns:
+        A tuple of two sets, where the first set contains the unique ops in the first program
+        and the second set contains the unique ops in the second program.
+    """
+    program_a_ops = set(_count_fx_targets(program_a))
+    program_b_ops = set(_count_fx_targets(program_b))
+    return program_a_ops - program_b_ops, program_b_ops - program_a_ops

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_building.py

@@ -0,0 +1,734 @@
+"""NOTES:
+
+We need a typing module that will handling Python to ONNX type promotion for use.
+For example, if we have torch.ops.aten.add(Tensor, 1.0), we need to promote 1.0
+to the same type as Tensor. The same thing needs to work for
+torch.ops.aten.add(1.0, Tensor) as well, which means we need a mechanism to`
+"""
+
+# mypy: allow-untyped-defs
+# mypy: disable-error-code=union-attr
+from __future__ import annotations
+
+import copy
+import inspect
+import logging
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, TYPE_CHECKING
+
+from onnxscript import evaluator
+
+import torch
+from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript
+from torch.onnx._internal.exporter import _errors, _schemas, _tensors
+
+
+if TYPE_CHECKING:
+    import onnx
+
+
+logger = logging.getLogger(__name__)
+
+ValidAttributeType = (
+    ir.TensorProtocol
+    | int
+    | float
+    | bool
+    | str
+    | Sequence[int]
+    | Sequence[float]
+    | None
+)
+
+AllowedArgType = ir.Value | Sequence[ir.Value | ValidAttributeType] | ValidAttributeType
+
+
+# Logic for adapting inputs from general Python or PyTorch inputs to ONNX ir.Value
+def _construct_named_inputs_and_attrs(
+    signature: ir.schemas.OpSignature,
+    args: Sequence[AllowedArgType],
+    kwargs: Mapping[str, AllowedArgType],
+) -> tuple[dict[str, AllowedArgType], dict[str, ValidAttributeType]]:
+    """Construct two mappings: name to inputs and named to attributes based on the signature and args/kwargs.
+
+    This function uses the OpSignature to determine which argument in args and kwargs corresponds to
+    which parameter in the signature. ONNX node inputs are stored in named_inputs, and attributes are
+    stored in named_attrs. If an _optional input_ is not provided, it is filled with None.
+
+    Args:
+        signature: The OpSignature for the node.
+        args: The positional arguments for the node.
+        kwargs: The keyword arguments for the node.
+
+    Returns:
+        A tuple of two mappings: named_inputs and named_attrs.
+
+    Raises:
+        ValueError: If a required parameter is not provided.
+    """
+    # 1. Construct the (named_inputs, named_attrs) mapping based on (args, kwargs) and the signature.
+    #   a. Loop over all parameters in the signature and args together
+    #   b. Depending on param.is_input, Record named_inputs[param.name] = arg or named_attrs[param.name] = arg
+    #   c. Handle kwargs as well
+    #   d. Fill in None if the input is not provided
+    named_inputs: dict[str, Any] = {}
+    named_attrs: dict[str, Any] = {}
+    reversed_args_stack = list(reversed(args))
+    for param in signature.params:
+        if isinstance(param, ir.schemas.Parameter):
+            # Handle inputs
+            if reversed_args_stack:
+                # First exhaust the positional arguments
+                if param.variadic:
+                    # Handle variadic arguments
+                    named_inputs[param.name] = tuple(args)
+                    reversed_args_stack.clear()
+                else:
+                    named_inputs[param.name] = reversed_args_stack.pop()  # type: ignore[assignment]
+            elif param.name in kwargs:
+                named_inputs[param.name] = kwargs[param.name]  # type: ignore[assignment]
+            elif param.required:
+                raise ValueError(
+                    f"Required parameter '{param.name}' is not provided. "
+                    f"Signature: {signature}. Args: {args}. Kwargs: {kwargs}."
+                )
+            else:
+                logger.debug(
+                    "Optional parameter '%s' is not provided. Added as None. Signature: %s",
+                    param.name,
+                    signature,
+                )
+                named_inputs[param.name] = None  # type: ignore[assignment]
+        else:
+            # Handle attributes
+            attribute: ValidAttributeType | ir.Attr
+            if not isinstance(param, ir.schemas.AttributeParameter):
+                raise AssertionError(f"Expected AttributeParameter, got {type(param)}")
+            if reversed_args_stack:
+                # First exhaust the positional arguments
+                attribute = reversed_args_stack.pop()  # type: ignore[assignment]
+            elif param.name in kwargs:
+                attribute = kwargs[param.name]  # type: ignore[assignment]
+            elif param.default is not None:
+                attribute = param.default
+            else:
+                attribute = None
+
+            if attribute is None:
+                if param.required:
+                    raise ValueError(
+                        f"Required attribute '{param.name}' is not provided. "
+                        f"Signature: {signature}. Args: {args}. Kwargs: {kwargs}."
+                    )
+                else:
+                    logger.debug(
+                        "Optional attribute '%s' is None. Dropped. Signature: %s",
+                        param.name,
+                        signature,
+                    )
+                    continue
+
+            if isinstance(attribute, ir.Attr):
+                # Turn the attribute from an default value into an actual parameter for the node
+                attr_copied = copy.copy(attribute)
+                # Make sure the name is the same as the parameter name and not the name of the default parameter
+                attr_copied.name = param.name
+                attribute = attr_copied
+
+            if isinstance(attribute, int) and param.type == ir.AttributeType.FLOAT:
+                # Convert the attribute to float if needed. This happens in PyTorch
+                # where an attribute marked as float can be passed as an int.
+                attribute = float(attribute)
+            named_attrs[param.name] = attribute
+    return named_inputs, named_attrs  # type: ignore[return-value]
+
+
+def _resolve_parameter_dtypes(
+    signature: ir.schemas.OpSignature, named_inputs: Mapping[str, AllowedArgType]
+) -> Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol]:
+    """Determine which parameter takes which type.
+
+    Handle non-tensor input corner cases and type promotion.
+
+    Requires:
+        All ir.Value in name_inputs should have type set. Their type should be
+        compatible with the type_constraint of the corresponding parameter in the signature.
+
+    Args:
+        signature: The OpSignature for the node.
+        named_inputs: The mapping of parameter names to their arguments.
+
+    Returns:
+        A mapping of Constraint names to ir.TypeProtocol.
+    """
+    #   a. Create type_binding: dict[str, ir.TypeProtocol]
+    #   b. Iterate over all named_inputs
+    #   b0. Find the corresponding parameter in the signature
+    #   b1. If the argument is a Python constant, skip.
+    #   b2. If the argument is a ir.Value, Bind {constraint: arg.type}.
+    type_binding = {}
+    for name, arg in named_inputs.items():
+        param = signature.params_map[name]
+        if not isinstance(param, ir.schemas.Parameter):
+            raise AssertionError(f"Expected Parameter, got {type(param)}")
+        if isinstance(arg, (int, float, bool, str, Sequence, torch.Tensor)):
+            # Skip the Python constants because we do not know what dtype they should take yet
+            continue
+        elif isinstance(arg, ir.Value):
+            if arg.type is None:
+                # Skip the ir.Value if the type is not set
+                continue
+            # NOTE: We assume arg.type is compatible with the type_constraint
+            if arg.type is None:
+                raise AssertionError(f"Expected type to be set for {arg}")
+            # TODO(justinchuby): Implement type promotion logic here.
+            type_binding[param.type_constraint] = arg.type
+    return type_binding
+
+
+def _determine_input_dtype(
+    param: ir.schemas.Parameter,
+    arg: AllowedArgType,
+    type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol],
+) -> ir.DataType:
+    """Determine the dtype of the input that is a mix of Python constants and ir.Value."""
+    if param.type_constraint in type_binding:
+        # A known dtype is available because it was resolved
+        return type_binding[param.type_constraint].dtype
+    if len(param.type_constraint.allowed_types) == 1:
+        # Only one type is allowed by the type constraint
+        return next(iter(param.type_constraint.allowed_types)).dtype
+
+    # No dtype information available. Infer from the Python constant or (in the Sequence case)
+    # from a mix of Python constants and ir.Value
+    if isinstance(arg, bool):
+        return ir.DataType.BOOL
+    if isinstance(arg, float):
+        return ir.DataType.FLOAT
+    if isinstance(arg, int):
+        return ir.DataType.INT64
+    if isinstance(arg, str):
+        return ir.DataType.STRING
+    if isinstance(arg, (ir.Tensor, ir.TensorProtocol)):
+        return arg.dtype
+    if isinstance(arg, complex):
+        return ir.DataType.FLOAT
+    if arg is None:
+        return ir.DataType.UNDEFINED
+
+    # Handle sequences
+    if isinstance(arg, (tuple, list)):
+        if len(arg) == 0:
+            # Special case: Treat empty sequence as INT64 as they are typically used for shape
+            return ir.DataType.INT64
+
+        # Try to obtain the dtype from one of the values
+        for val in arg:
+            if isinstance(val, ir.Value) and val.dtype is not None:
+                return val.dtype
+
+        if any(isinstance(val, float) for val in arg):
+            # If any float is present, the dtype is float
+            return ir.DataType.FLOAT
+        elif any(isinstance(val, int) for val in arg):
+            # Otherwise if any int is present, the dtype is int
+            return ir.DataType.INT64
+
+    raise ValueError(
+        f"Could not determine the dtype for the input '{param.name}'. "
+        f"param={param}, arg={arg}, param_type_constraint={param.type_constraint}, "
+        f"type_binding={type_binding}"
+    )
+
+
+def _allowed_types_are_sequence_types(allowed_types: Iterable[ir.TypeProtocol]) -> bool:
+    """Check if all allowed types are Sequence types."""
+    return all(isinstance(t, ir.SequenceType) for t in allowed_types)
+
+
+def _get_or_create_constant(
+    constant_farm: dict[
+        tuple[
+            bool
+            | int
+            | float
+            | str
+            | tuple[int, ...]
+            | tuple[float, ...]
+            | tuple[bool, ...],
+            ir.DataType,
+        ],
+        ir.Value,
+    ],
+    arg: bool
+    | int
+    | float
+    | str
+    | tuple[int, ...]
+    | tuple[float, ...]
+    | tuple[bool, ...]
+    | list[int]
+    | list[float]
+    | list[bool],
+    dtype: ir.DataType,
+    opset: onnxscript.values.Opset,
+) -> ir.Value:
+    # float representation of complex numbers
+    if isinstance(arg, complex):
+        # Convert the complex number to a float
+        arg = (arg.real, arg.imag)
+
+    if isinstance(arg, list):
+        # Make the arg hashable
+        # pyrefly: ignore [bad-argument-type]
+        arg = tuple(arg)
+
+    constant_value = constant_farm.get((arg, dtype))  # type: ignore[arg-type]
+    if constant_value is None:
+        constant_tensor = ir.tensor(value=arg, dtype=dtype)
+        constant_value = opset.Constant(value=constant_tensor)
+        constant_farm[(arg, dtype)] = constant_value  # type: ignore[arg-type,index]
+    return constant_value  # type: ignore[return-value]
+
+
+def _process_python_constants(
+    signature: ir.schemas.OpSignature,
+    named_inputs: dict[str, AllowedArgType],
+    type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol],
+    constant_farm: dict[
+        tuple[
+            bool | int | float | str | tuple[int, ...] | tuple[float, ...],
+            ir.DataType,
+        ],
+        ir.Value,
+    ],
+    opset: onnxscript.values.Opset,
+) -> dict[str, ir.Value | None]:
+    """Convert Python constants to Constant nodes and list to Sequence nodes based on the dtype information.
+
+    The added constants will be replacing values in named_inputs in place.
+
+    Args:
+        signature: The OpSignature for the node.
+        named_inputs: The mapping of parameter names to their arguments.
+        type_binding: A mapping of Constraint names to ir.DataType.
+        constant_farm: A dictionary of {(py_value, ir.DataType): ir.Value} to store the deduplicated constants.
+        opset: The Opset to use for creating Constant nodes.
+
+    Returns:
+        A mapping of parameter names to Python constants converted to constant Nodes.
+    """
+    # 3. Convert Python constants to Constant nodes based on the dtype information;
+    #    construct sequences
+    #   a. Iterate over all parameters in the signature the second time
+    #   b. If the parameter is in to_resolve_type:
+    #       - If param.constraint in type_binding,
+    #         Get the constant from constant_farm (deduplicated);
+    #            otherwise set named_inputs[param.name] = Constant(value, dtype=type_binding[param.constraint])
+    #       - Otherwise, set named_inputs[param.name] = Constant(value)
+    for name, arg in named_inputs.items():
+        param = signature.params_map[name]
+        if not isinstance(param, ir.schemas.Parameter):
+            raise AssertionError(f"Expected Parameter, got {type(param)}")
+
+        if isinstance(arg, ir.Value):
+            # TODO(justinchuby): Cast the ir.Value here if needed
+            continue
+
+        if (
+            isinstance(arg, Sequence)
+            and len(arg) > 0
+            and any(isinstance(val, ir.Value) for val in arg)
+        ):
+            # Skip the sequence of ir.Value. This is a variadic input or a Sequence input
+            # It will be handled by _process_python_sequences
+            continue
+        if param.variadic:
+            # Handled by _process_python_sequences
+            continue
+        if _allowed_types_are_sequence_types(param.type_constraint.allowed_types):
+            # Handled by _process_python_sequences
+            continue
+
+        dtype = _determine_input_dtype(param, arg, type_binding)
+
+        if arg is None:
+            constant_value = None
+        elif isinstance(arg, (ir.Tensor, ir.TensorProtocol)):
+            constant_value = opset.Constant(value=arg)
+        else:
+            # Deduplicate the constants
+            constant_value = _get_or_create_constant(constant_farm, arg, dtype, opset)  # type: ignore[arg-type]
+
+        named_inputs[param.name] = constant_value
+    return named_inputs  # type: ignore[return-value]
+
+
+def _reshape_to_1d_tensor(opset: onnxscript.values.Opset, arg: ir.Value) -> ir.Value:
+    """Reshape the input to a 1D tensor."""
+
+    return opset.Reshape(
+        arg, opset.Constant(value=ir.tensor([-1], dtype=ir.DataType.INT64))
+    )
+
+
+def _process_python_sequences(
+    signature: ir.schemas.OpSignature,
+    named_inputs: dict[str, AllowedArgType],
+    type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol],
+    constant_farm: dict[
+        tuple[
+            bool
+            | int
+            | float
+            | str
+            | ir.TensorProtocol
+            | tuple[bool, ...]
+            | tuple[int, ...]
+            | tuple[float, ...],
+            ir.DataType,
+        ],
+        ir.Value,
+    ],
+    opset: onnxscript.values.Opset,
+):
+    """Handle three types of sequences.
+
+    1. Variadic inputs
+    2. Sequence input of ir.Value,
+    3. Sequence of Python constants that contains ir.Value
+    """
+    for name, arg in named_inputs.items():
+        param = signature.params_map[name]
+        if not isinstance(param, ir.schemas.Parameter):
+            raise AssertionError(f"Expected Parameter, got {type(param)}")
+
+        if not isinstance(arg, (tuple, list)):
+            continue
+
+        if len(arg) == 0:
+            # Skip empty sequences
+            continue
+
+        # 1. Sequence input of ir.Value
+        if _allowed_types_are_sequence_types(param.type_constraint.allowed_types):
+            # Turn the list into a Sequence node
+            # Constant op creation will be handled by the variadic case below when calling
+            # the SequenceConstruct op.
+            named_inputs[name] = opset.SequenceConstruct(*arg)
+            continue
+
+        # 2. Variadic inputs
+        # NOTE: Variadic operators like Max can be called with mixed ir.Value and Python constants
+        # like `Max(0, ir.Value())`
+        # We need to convert the Python constants to Constant nodes
+        if param.variadic:
+            if all(isinstance(val, ir.Value) for val in arg):
+                # Skip the variadic input if all values are ir.Value
+                continue
+
+            dtype = _determine_input_dtype(param, arg, type_binding)
+            new_args = []
+            for val in arg:
+                if isinstance(val, ir.Value):
+                    new_args.append(val)
+                else:
+                    constant_tensor = ir.tensor(value=val, dtype=dtype)  # type: ignore[arg-type]
+                    constant_value = opset.Constant(value=constant_tensor)
+                    new_args.append(constant_value)
+            named_inputs[name] = new_args
+            continue
+        else:
+            # 3. Concat the list as a single input
+            # E.g. [Value, 42] should be converted to op.Concat(Value, Constant(42))
+            # when the expected input type is INT64
+            # We assume this only happens for 0D cases
+            if all(isinstance(val, ir.Value) for val in arg):
+                expanded_args = [_reshape_to_1d_tensor(opset, val) for val in arg]
+                named_inputs[name] = opset.Concat(*expanded_args, axis=0)
+                continue
+
+            dtype = _determine_input_dtype(param, arg, type_binding)
+            new_args = []
+            for val in arg:
+                if isinstance(val, ir.Value):
+                    new_args.append(_reshape_to_1d_tensor(opset, val))
+                elif val is None:
+                    # Skip None values
+                    continue
+                elif isinstance(val, (ir.Tensor, ir.TensorProtocol)):
+                    new_args.append(
+                        _reshape_to_1d_tensor(opset, opset.Constant(value=val))
+                    )
+                else:
+                    # Turn the Python constant into 1D tensor for the constant
+                    if not isinstance(val, (bool, int, float)):
+                        raise AssertionError(f"Expected int or float, got {type(val)}")
+                    new_args.append(
+                        _get_or_create_constant(constant_farm, [val], dtype, opset)  # type: ignore[arg-type]
+                    )
+            named_inputs[name] = opset.Concat(*new_args, axis=0)
+            continue
+    return named_inputs
+
+
+def _determine_output_number(
+    signature: ir.schemas.OpSignature, named_attrs: Mapping[str, ValidAttributeType]
+) -> int:
+    """Determine the number of outputs for the node with heuristics."""
+    if signature.domain == "":
+        if signature.name == "BatchNormalization":
+            if not named_attrs.get("training_mode", 0):
+                return 1
+        if signature.name == "Split":
+            num_outputs = named_attrs.get("num_outputs")
+            if num_outputs is not None and isinstance(num_outputs, int):
+                return num_outputs
+            else:
+                raise ValueError(
+                    "Could not determine the number of outputs for Split. "
+                    "num_outputs must be provided"
+                )
+    return len(signature.outputs)
+
+
+def _construct_node(
+    signature: ir.schemas.OpSignature,
+    named_inputs: Mapping[str, ir.Value | None],
+    named_attrs: Mapping[str, ValidAttributeType],
+    opset: onnxscript.values.Opset,
+    num_outputs: int,
+) -> ir.Node:
+    """Construct the node with the inputs and attributes.
+
+    Variadic inputs are flattened.
+
+    Args:
+        signature: The OpSignature for the node.
+        named_inputs: The mapping of parameter names to their arguments. When we
+            do not have the schema of an operator, we do not know the names of
+            the inputs, in which case the names can be anything because they
+            are not used in this function. The data structure is passed in for
+            consistency with the other functions.
+        named_attrs: The mapping of attribute names to their values.
+        num_outputs: The number of outputs for the node.
+    """
+    inputs: list[ir.Value | None] = []
+    # Flatten variadic inputs
+    for value in named_inputs.values():
+        if isinstance(value, Sequence):
+            inputs.extend(value)
+        else:
+            inputs.append(value)
+
+    # If final inputs are None, strip them from the node inputs
+    for input in reversed(inputs):
+        if input is not None:
+            break
+        inputs.pop()
+
+    # Construct and filter out None attributes
+    attributes = [
+        attr
+        for attr in ir.convenience.convert_attributes(named_attrs)
+        if attr.value is not None
+    ]
+    outputs = [_tensors.SymbolicTensor(opset) for _ in range(num_outputs)]
+    return ir.Node(
+        signature.domain,
+        signature.name,
+        inputs=inputs,
+        attributes=attributes,
+        outputs=outputs,
+        version=signature.since_version,
+    )
+
+
+class OpRecorder(evaluator.Evaluator):
+    """An onnxscript Evaluator that captures the graph into ONNX IR."""
+
+    def __init__(
+        self, opset: onnxscript.values.Opset, constant_farm: dict[Any, ir.Value]
+    ) -> None:
+        self.nodes: list[ir.Node] = []
+        self.opset = opset
+        self.functions: dict[
+            ir.OperatorIdentifier, onnxscript.OnnxFunction | ir.Function
+        ] = {}
+        self.constant_farm = constant_farm
+
+    def _call_op(
+        self,
+        op_signature: ir.schemas.OpSignature,
+        named_inputs: dict[str, AllowedArgType],
+        named_attrs: dict[str, ValidAttributeType],
+        num_outputs: int,
+    ) -> Sequence[_tensors.SymbolicTensor]:
+        """Record nodes for the given opschema and arguments.
+
+        Args:
+            op_signature: The OpSchema containing the node signature.
+            named_inputs: The mapping of parameter names to their arguments.
+            named_attrs: The mapping of attribute names to their values.
+        """
+        type_binding = _resolve_parameter_dtypes(op_signature, named_inputs)
+        try:
+            converted_named_inputs = _process_python_constants(
+                op_signature, named_inputs, type_binding, self.constant_farm, self.opset
+            )
+            converted_named_inputs = _process_python_sequences(
+                op_signature,
+                converted_named_inputs,  # type: ignore[arg-type]
+                type_binding,
+                self.constant_farm,
+                self.opset,
+            )
+
+        except Exception as e:
+            raise _errors.GraphConstructionError(
+                f"Error processing Python constants for operator '{op_signature.domain}::{op_signature.name}'. "
+                f"named_inputs={named_inputs}, named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}."
+            ) from e
+
+        try:
+            self.nodes.append(
+                node := _construct_node(
+                    op_signature,
+                    converted_named_inputs,
+                    named_attrs,
+                    self.opset,
+                    num_outputs,
+                )
+            )
+        except Exception as e:
+            raise _errors.GraphConstructionError(
+                f"Error constructing node for operator '{op_signature.domain}::{op_signature.name}'. "
+                f"named_inputs={named_inputs}, converted_named_inputs={converted_named_inputs}, "
+                f"named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}."
+            ) from e
+        return node.outputs  # type: ignore[return-value]
+
+    def eval(
+        self,
+        schema: onnx.defs.OpSchema,
+        args: Sequence[AllowedArgType],  # type: ignore[override]
+        kwargs: Mapping[str, AllowedArgType],
+    ) -> _tensors.SymbolicTensor | Sequence[_tensors.SymbolicTensor]:
+        try:
+            op_signature = ir.schemas.OpSignature.from_op_schema(schema)
+            named_inputs, named_attrs = _construct_named_inputs_and_attrs(
+                op_signature, args, kwargs
+            )
+            # TODO(justinchuby): Handle cast
+            if schema.name == "CastLike":
+                if len(named_inputs) != 2:
+                    raise AssertionError(f"Expected 2 inputs, got {len(named_inputs)}")
+                # Skip CastLike if the input and output types are the same
+                src_input = named_inputs["input"]
+                target_type = named_inputs["target_type"]
+
+                if (
+                    isinstance(src_input, ir.Value)
+                    and isinstance(target_type, ir.Value)
+                    and src_input.dtype is not None
+                    and target_type.dtype is not None
+                ):
+                    # dtypes are available
+                    if src_input.dtype == target_type.dtype:
+                        # Same type. No cast needed
+                        return src_input  # type: ignore[return-value]
+                    else:
+                        # Create a Cast node
+                        return self.opset.Cast(src_input, to=target_type.dtype)  # type: ignore[union-attr,return-value]
+
+            num_outputs = _determine_output_number(op_signature, named_attrs)
+            outputs = self._call_op(
+                op_signature, named_inputs, named_attrs, num_outputs
+            )
+            if len(outputs) == 1:
+                return outputs[0]
+            return outputs
+        except Exception as e:
+            raise _errors.GraphConstructionError(
+                f"Error calling operator '{schema.name}' with args {args} and kwargs {kwargs}."
+            ) from e
+
+    def eval_function(  # type: ignore[override]
+        self,
+        function: onnxscript.OnnxFunction,
+        args: Sequence[AllowedArgType],
+        kwargs: Mapping[str, AllowedArgType],
+    ) -> _tensors.SymbolicTensor | Sequence[_tensors.SymbolicTensor] | bool | int:
+        try:
+            # NOTE: signature should be written to function in the registration process
+            if hasattr(function, "_pt_onnx_signature"):
+                op_signature = function._pt_onnx_signature  # type: ignore[attr-defined]
+            else:
+                op_signature = _schemas.op_signature_from_function(
+                    function,
+                    function.function_ir.domain,
+                    function.name,
+                    since_version=function.opset.version,
+                )
+                function._pt_onnx_signature = op_signature  # type: ignore[attr-defined]
+
+            named_inputs, named_attrs = _construct_named_inputs_and_attrs(
+                op_signature, args, kwargs
+            )
+
+            # TODO(after torchlib migration): Remove traceable function handling
+            # NOTE: We need to call traceable functions after the _construct_named_inputs_and_attrs
+            # call because it will filter out the unexpected kwargs for us.
+            if function.traceable:
+                # Trace the function call instead of adding the function as a node
+                # Turn the ir.Attr objects into Python constants first
+                named_attrs = {
+                    name: attr.value if isinstance(attr, ir.Attr) else attr
+                    for name, attr in named_attrs.items()
+                }
+
+                # Use the type binding to resolve the dtypes of the inputs, and
+                # convert Python constants to Constant nodes
+                type_binding = _resolve_parameter_dtypes(op_signature, named_inputs)
+                try:
+                    # _process_python_sequences is not here because we want to preserve python list
+                    # properties for the function call
+                    converted_named_inputs = _process_python_constants(
+                        op_signature,
+                        named_inputs,
+                        type_binding,
+                        self.constant_farm,
+                        self.opset,
+                    )
+
+                except Exception as e:
+                    raise _errors.GraphConstructionError(
+                        f"Error processing Python constants for operator '{op_signature.domain}::{op_signature.name}'. "
+                        f"named_inputs={named_inputs}, named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}."
+                    ) from e
+
+                return function.function(**converted_named_inputs, **named_attrs)
+
+            outputs = self._call_op(
+                op_signature,
+                named_inputs,
+                named_attrs,
+                len(op_signature.outputs),
+            )
+
+            self.functions[(function.function_ir.domain, function.name, "")] = function
+            if len(outputs) == 1:
+                return outputs[0]
+            return outputs
+        except Exception as e:
+            try:
+                source_file = inspect.getsourcefile(function.function)
+                _, lineno = inspect.getsourcelines(function.function)
+            except Exception:
+                source_file = lineno = None
+            raise _errors.GraphConstructionError(
+                f"Error calling function '{function.name}' with args {args} and kwargs {kwargs}."
+                + f" The function is defined at '{source_file}:{lineno}'."
+                if source_file
+                else ""
+            ) from e

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py

@@ -0,0 +1,304 @@
+"""Strategies for capturing ExportedPrograms."""
+
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import abc
+import contextlib
+import dataclasses
+import datetime
+import logging
+import pathlib
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.onnx import _flags
+
+
+if TYPE_CHECKING:
+    import os
+    from collections.abc import Callable
+
+
+logger = logging.getLogger(__name__)
+
+
+def _verbose_printer(verbose: bool | None) -> Callable[..., None]:
+    """Prints messages based on `verbose`."""
+    if verbose is False:
+        return lambda *_, **__: None
+
+    return lambda *args, **kwargs: print("[torch.onnx]", *args, **kwargs)
+
+
+def _take_first_line(text: str) -> str:
+    """Take the first line of a text."""
+    lines = text.split("\n", maxsplit=1)
+    first_line = lines[0]
+    if len(lines) > 1:
+        first_line += "[...]"
+    return first_line
+
+
+@contextlib.contextmanager
+def _patch_dynamo_unsupported_functions():
+    """Patch PyTorch to bypass some functions torch.export.export does not support."""
+    # TODO: Remove the patches once dynamo supports these functions.
+    import torch.jit
+
+    # Replace torch.jit.isinstance with isinstance
+    jit_isinstance = torch.jit.isinstance
+    # pyrefly: ignore [bad-assignment]
+    torch.jit.isinstance = isinstance
+    logger.info("Replaced torch.jit.isinstance with isinstance to allow dynamo tracing")
+    try:
+        yield
+    finally:
+        torch.jit.isinstance = jit_isinstance
+
+
+@dataclasses.dataclass
+class Result:
+    exported_program: torch.export.ExportedProgram | None
+    strategy: str
+    exception: Exception | None = None
+
+    @property
+    def success(self) -> bool:
+        """Whether the capture was successful.
+
+        An exception can still be recorded even if the capture was successful. In
+        this case the exception is informational only. For example, draft_export
+        can record an exception if there are warnings during the export. The exceptions
+        will go into the onnx export report when report=True.
+        """
+        return self.exported_program is not None
+
+
+class CaptureStrategy(abc.ABC):
+    """Strategy for capturing a module as ExportedProgram.
+
+    To use a strategy, create an instance and call it with the model, args, kwargs, and dynamic_shapes.
+    Example::
+
+        strategy = TorchExportNonStrictStrategy(verbose=True)
+        result = strategy(model, args, kwargs, dynamic_shapes)
+    """
+
+    def __init__(
+        self,
+        *,
+        verbose: bool = False,
+        dump: bool = False,
+        artifacts_dir: str | os.PathLike = ".",
+        timestamp: str | None = None,
+    ) -> None:
+        """Initialize the strategy.
+
+        Args:
+            verbose: Whether to print verbose messages.
+            dump: Whether to dump the intermediate artifacts to a file.
+        """
+        self._verbose_print = _verbose_printer(verbose)
+        self._dump = dump
+        self._artifacts_dir = pathlib.Path(artifacts_dir)
+        self._timestamp = timestamp or datetime.datetime.now().strftime(
+            "%Y-%m-%d_%H-%M-%S-%f"
+        )
+        self._exception: Exception | None = None
+
+    def __call__(
+        self,
+        model: torch.nn.Module | torch.jit.ScriptFunction,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any] | None,
+        dynamic_shapes,
+    ) -> Result:
+        self._enter(model)
+        if kwargs is None:
+            kwargs = {}
+        try:
+            exported_program = self._capture(model, args, kwargs, dynamic_shapes)
+        except Exception as e:
+            self._failure(model, e)
+            return Result(
+                exported_program=None,
+                strategy=self.__class__.__name__,
+                exception=e,
+            )
+        self._success(model)
+        return Result(
+            exported_program,
+            strategy=self.__class__.__name__,
+            exception=self._exception,
+        )
+
+    @abc.abstractmethod
+    def _capture(
+        self, model, args, kwargs, dynamic_shapes
+    ) -> torch.export.ExportedProgram:
+        raise NotImplementedError
+
+    def _enter(self, model: torch.nn.Module | torch.jit.ScriptFunction) -> None:
+        return
+
+    def _success(self, model: torch.nn.Module | torch.jit.ScriptFunction) -> None:
+        return
+
+    def _failure(
+        self, model: torch.nn.Module | torch.jit.ScriptFunction, e: Exception
+    ) -> None:
+        return
+
+
+class TorchExportStrictStrategy(CaptureStrategy):
+    def _capture(
+        self, model, args, kwargs, dynamic_shapes
+    ) -> torch.export.ExportedProgram:
+        with (
+            _patch_dynamo_unsupported_functions(),
+            # Support the dynamism with 0/1 input dim
+            torch.fx.experimental._config.patch(backed_size_oblivious=True),  # type: ignore[attr-defined]
+        ):
+            try:
+                return torch.export.export(
+                    model,
+                    args,
+                    kwargs=kwargs,
+                    dynamic_shapes=dynamic_shapes,
+                    strict=True,
+                    prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS,
+                )
+            except torch._dynamo.exc.UserError as exc:
+                # Refine the dynamic shapes based on the suggested fixes.
+                try:
+                    new_shapes = torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(
+                        exc.msg, dynamic_shapes
+                    )
+                except Exception:
+                    # If the dynamic shapes cannot be refined, re-raise the exception.
+                    raise exc from None
+                return torch.export.export(
+                    model,
+                    args,
+                    kwargs=kwargs,
+                    dynamic_shapes=new_shapes,
+                    strict=True,
+                    prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS,
+                )
+
+    def _enter(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`..."
+        )
+
+    def _success(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`... ✅"
+        )
+
+    def _failure(self, model, e) -> None:
+        del e  # Unused
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`... ❌"
+        )
+
+
+class TorchExportNonStrictStrategy(CaptureStrategy):
+    def _capture(
+        self, model, args, kwargs, dynamic_shapes
+    ) -> torch.export.ExportedProgram:
+        with (
+            # Support the dynamism with 0/1 input dim
+            torch.fx.experimental._config.patch(backed_size_oblivious=True),  # type: ignore[attr-defined]
+        ):
+            try:
+                return torch.export.export(
+                    model,
+                    args,
+                    kwargs=kwargs,
+                    dynamic_shapes=dynamic_shapes,
+                    strict=False,
+                    prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS,
+                )
+            except torch._dynamo.exc.UserError as exc:
+                # Refine the dynamic shapes based on the suggested fixes.
+                try:
+                    new_shapes = torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(
+                        exc.msg, dynamic_shapes
+                    )
+                except Exception:
+                    # If the dynamic shapes cannot be refined, re-raise the exception.
+                    raise exc from None
+                return torch.export.export(
+                    model,
+                    args,
+                    kwargs=kwargs,
+                    dynamic_shapes=new_shapes,
+                    strict=False,
+                    prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS,
+                )
+
+    def _enter(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`..."
+        )
+
+    def _success(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`... ✅"
+        )
+
+    def _failure(self, model, e) -> None:
+        del e  # Unused
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`... ❌"
+        )
+
+
+class TorchExportDraftExportStrategy(CaptureStrategy):
+    def _capture(
+        self, model, args, kwargs, dynamic_shapes
+    ) -> torch.export.ExportedProgram:
+        ep = torch.export.draft_export(
+            model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes
+        )
+        report = ep._report  # type: ignore[attr-defined]
+        if not report.successful():
+            self._exception = RuntimeError(str(report))
+            self._verbose_print(f"Draft Export report:\n{report}")
+        return ep
+
+    def _enter(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`..."
+        )
+
+    def _success(self, model) -> None:
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`... ✅"
+        )
+
+    def _failure(self, model, e) -> None:
+        del e  # Unused
+        model_repr = _take_first_line(repr(model))
+        self._verbose_print(
+            f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`... ❌"
+        )
+
+
+CAPTURE_STRATEGIES: tuple[type[CaptureStrategy], ...] = (
+    TorchExportNonStrictStrategy,  # strict=False is preferred over strict=True because it does not have dynamo issues
+    TorchExportStrictStrategy,
+)
+
+if _flags.ENABLE_DRAFT_EXPORT:
+    CAPTURE_STRATEGIES = (*CAPTURE_STRATEGIES, TorchExportDraftExportStrategy)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_compat.py

@@ -0,0 +1,202 @@
+"""Compatibility functions for the torch.onnx.export API."""
+
+# mypy: allow-untyped-defs
+# mypy: disable-error-code=attr-defined
+from __future__ import annotations
+
+import io
+import logging
+import warnings
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.onnx import _constants as onnx_constants
+from torch.onnx._internal._lazy_import import onnx
+from torch.onnx._internal.exporter import (
+    _constants,
+    _core,
+    _dynamic_shapes,
+    _exportable_module,
+    _onnx_program,
+    _registration,
+)
+
+
+if TYPE_CHECKING:
+    import os
+
+logger = logging.getLogger(__name__)
+
+
+def _get_torch_export_args(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+) -> tuple[tuple[Any, ...], dict[str, Any] | None]:
+    """Obtain the arguments for torch.onnx.export from the model and the input arguments."""
+    if not kwargs and args and isinstance(args[-1], dict):
+        kwargs = args[-1]
+        args = args[:-1]
+    return args, kwargs
+
+
+def export_compat(
+    model: torch.nn.Module
+    | torch.export.ExportedProgram
+    | torch.jit.ScriptModule
+    | torch.jit.ScriptFunction,
+    args: tuple[Any, ...],
+    f: str | os.PathLike | None = None,
+    *,
+    kwargs: dict[str, Any] | None = None,
+    export_params: bool = True,
+    verbose: bool | None = None,
+    input_names: Sequence[str] | None = None,
+    output_names: Sequence[str] | None = None,
+    opset_version: int | None = onnx_constants.ONNX_DEFAULT_OPSET,
+    custom_translation_table: dict[Callable, Callable] | None = None,
+    dynamic_axes: Mapping[str, Mapping[int, str]]
+    | Mapping[str, Sequence[int]]
+    | None = None,
+    dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,
+    keep_initializers_as_inputs: bool = False,
+    external_data: bool = True,
+    report: bool = False,
+    optimize: bool = True,
+    verify: bool = False,
+    profile: bool = False,
+    dump_exported_program: bool = False,
+    artifacts_dir: str | os.PathLike = ".",
+) -> _onnx_program.ONNXProgram:
+    if opset_version is None:
+        opset_version = onnx_constants.ONNX_DEFAULT_OPSET
+
+    if isinstance(model, torch.nn.Module):
+        if model.training:
+            warnings.warn(
+                "Exporting a model while it is in training mode. "
+                "Please ensure that this is intended, as it may lead to "
+                "different behavior during inference. "
+                "Calling model.eval() before export is recommended.",
+                UserWarning,
+                stacklevel=3,
+            )
+
+    if isinstance(model, _exportable_module.ExportableModule):
+        # Skip argument extraction if args or kwargs are provided
+        if not args and not kwargs:
+            args, kwargs = model.example_arguments()
+            if input_names is None:
+                input_names = model.input_names()
+            if output_names is None:
+                output_names = model.output_names()
+            if dynamic_shapes is None:
+                dynamic_shapes = model.dynamic_shapes()
+
+    if isinstance(model, torch.export.ExportedProgram):
+        # We know the model is already exported program, so the args, kwargs, and dynamic_shapes
+        # are not used
+        dynamic_shapes = dynamic_shapes or {}
+    else:
+        args, kwargs = _get_torch_export_args(args, kwargs)
+        if dynamic_shapes is None and dynamic_axes is not None:
+            warnings.warn(
+                "# 'dynamic_axes' is not recommended when dynamo=True, "
+                "and may lead to 'torch._dynamo.exc.UserError: Constraints violated.' "
+                "Supply the 'dynamic_shapes' argument instead if export is unsuccessful.",
+                UserWarning,
+                stacklevel=3,
+            )
+            try:
+                dynamic_shapes, args, kwargs = (
+                    _dynamic_shapes.from_dynamic_axes_to_dynamic_shapes(
+                        model,
+                        args,
+                        kwargs,
+                        dynamic_axes=dynamic_axes,
+                        input_names=input_names,
+                        output_names=set(output_names or ()),
+                    )
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    "# Failed to convert 'dynamic_axes' to 'dynamic_shapes'. "
+                    "Please provide 'dynamic_shapes' directly. "
+                    "Refer to the documentation for 'torch.export.export' for more information on dynamic shapes."
+                ) from e
+
+    dynamic_shapes_with_export_dim, need_axis_mapping = (
+        _dynamic_shapes.convert_str_to_export_dim(dynamic_shapes)
+    )
+
+    if opset_version < _constants.TORCHLIB_OPSET:
+        logger.warning(
+            "Setting ONNX exporter to use operator set version %s because "
+            "the requested opset_version %s is a lower version than we have implementations for. "
+            "Automatic version conversion will be performed, which may not be successful "
+            "at converting to the requested version. If version conversion is unsuccessful, "
+            "the opset version of the exported model will be kept at %s. "
+            "Please consider setting opset_version >=%s to leverage latest ONNX features",
+            _constants.TORCHLIB_OPSET,
+            opset_version,
+            _constants.TORCHLIB_OPSET,
+            _constants.TORCHLIB_OPSET,
+        )
+        registry_opset_version = _constants.TORCHLIB_OPSET
+    else:
+        registry_opset_version = opset_version
+
+    registry = _registration.ONNXRegistry().from_torchlib(
+        opset_version=registry_opset_version
+    )
+    if custom_translation_table is not None:
+        for torch_op, onnx_op in custom_translation_table.items():
+            # TODO(justinchuby): Support complex inputs with annotations
+            if isinstance(onnx_op, Sequence):
+                raise TypeError(
+                    "The value in custom_translation_table should be a single callable, not a sequence"
+                )
+            registry.register_op(torch_op, onnx_op, is_complex=False)
+
+    onnx_program = _core.export(
+        model,
+        args,
+        kwargs,
+        registry=registry,
+        dynamic_shapes=dynamic_shapes_with_export_dim,
+        input_names=input_names,
+        output_names=output_names,
+        profile=profile,
+        report=report,
+        verify=verify,
+        dump_exported_program=dump_exported_program,
+        artifacts_dir=artifacts_dir,
+        verbose=verbose,
+        optimize=optimize,
+        opset_version=opset_version,
+    )
+
+    if need_axis_mapping and dynamic_shapes is not None:
+        onnx_program._rename_dynamic_axes(dynamic_shapes)
+
+    if f is not None:
+        if isinstance(f, io.BytesIO):
+            # For legacy export compatibility, we allow f to be a BytesIO object.
+            # This is not explicitly supported but we may need to maintain the
+            # behavior indefinitely.
+            warnings.warn(
+                "Saving ONNX model to a BytesIO object is deprecated. "
+                "Please use a file path instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            onnx.save(onnx_program.model_proto, f)
+        else:
+            onnx_program.save(
+                f,
+                include_initializers=export_params,
+                keep_initializers_as_inputs=keep_initializers_as_inputs,
+                external_data=external_data,
+            )
+
+    return onnx_program

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_constants.py

@@ -0,0 +1,7 @@
+# ir_version used for the ONNX file. See https://github.com/onnx/onnx/blob/main/docs/IR.md#onnx-versioning
+ONNX_IR_VERSION = 10
+# The opset version torchlib is implemented with. Update this number when updating torchlib
+TORCHLIB_OPSET = 18
+TORCHLIB_DOMAIN = "pkg.torch.onnx"
+# Domain used for functions translated from subgraphs
+LOCAL_FUNCTION_DOMAIN = "pkg.torch.__subgraph__"

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_decomp.py

@@ -0,0 +1,77 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import itertools
+from typing import TYPE_CHECKING
+
+import torch
+import torch._ops
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.onnx._internal.exporter import _registration
+
+
+def get_onnx_implemented_overloads(
+    registry: _registration.ONNXRegistry,
+) -> list[_registration.TorchOp]:
+    """
+    Creates a set of OperatorBase and Callable objects that represent ONNX-supported PyTorch operations.
+
+    Args:
+        registry: The ONNX registry for PyTorch.
+
+    Returns:
+        A collection of OperatorBase and Callable objects representing ONNX-supported PyTorch operations.
+    """
+    registered_ops: list[_registration.TorchOp] = []
+    for onnx_decomp_meta in registry.functions.values():
+        if len(onnx_decomp_meta) == 0:
+            raise AssertionError("onnx_decomp_meta must not be empty")
+        # Different OnnxDecompMeta for the same TorchOp should
+        # have the same fx_target.
+        fx_target = onnx_decomp_meta[0].fx_target
+        registered_ops.append(fx_target)
+    return registered_ops
+
+
+def create_onnx_friendly_decomposition_table(
+    onnx_registered_ops: set[_registration.TorchOp],
+) -> dict[_registration.TorchOp, Callable]:
+    """
+    This function creates a dictionary of op overloads and their decomposition functions
+    for ops that do not have ONNX symbolic functions. If an op already has an ONNX symbolic function,
+    its decomposition function is excluded from the table. The decomposition table is a subset of PyTorch's
+    built-in aten-to-aten decomposition.
+
+    Args:
+        onnx_registered_ops: All ops that have an ONNX decomposition implemented.
+
+    Returns:
+        Dict[torch._ops.OperatorBase, Callable]: A dictionary that maps op overloads to their corresponding
+        decomposition functions.
+    """
+    decomposition_table: dict[_registration.TorchOp, Callable] = {}
+
+    for op_overload, decomp_fn in itertools.chain(
+        torch.export.default_decompositions().items(),  # type: ignore[attr-defined]
+        torch._decomp.decomposition_table.items(),  # type: ignore[attr-defined]
+    ):
+        # Skip decomposition for op_overload as long as that op_overload has a corresponding ONNX
+        # symbolic function.
+        # NOTE: Do not skip torch._refs decomps. They are fine because otherwise the model is
+        # not exportable anyways.
+        if op_overload in onnx_registered_ops:
+            continue
+        # If it is HOP, we filter those out as well.
+        if not hasattr(op_overload, "_schema"):
+            continue
+        # NOTE: torch._decomp.decomposition_table covers more ops
+        # than torch.export.default_decompositions, but the latter is
+        # more critical to torch.onnx.export.
+        if op_overload in decomposition_table:
+            continue
+        decomposition_table[op_overload] = decomp_fn
+    return decomposition_table

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_dispatching.py

@@ -0,0 +1,58 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.fx
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.onnx._internal.exporter import _registration
+
+
+def _arg_has_complex_dtype(arg) -> bool:
+    """Check if the node has complex dtype recursively."""
+    if (
+        isinstance(arg, torch.fx.Node)
+        and "val" in arg.meta
+        and isinstance(arg.meta["val"], torch.Tensor)
+        and torch.is_complex(arg.meta["val"])
+    ):
+        return True
+    elif isinstance(arg, list):
+        return any(_arg_has_complex_dtype(item) for item in arg)
+    return False
+
+
+def dispatch(
+    node: torch.fx.Node, registry: _registration.ONNXRegistry
+) -> tuple[Callable | None, str]:
+    """Dispatch a node to an ONNX function based on the node's target and the ONNX registry.
+
+    Args:
+        node: The node to dispatch.
+        registry: The ONNX registry to use for dispatching.
+
+    Returns:
+        A tuple containing the matched ONNX function and a string describing the reason for failure or success.
+    """
+    decomp_metas = registry.get_decomps(node.target)  # type: ignore[arg-type]
+    # Determine if the node has complex inputs.
+    is_complex = any(_arg_has_complex_dtype(arg) for arg in node.args) or any(
+        _arg_has_complex_dtype(arg) for arg in node.kwargs.values()
+    )
+    if is_complex:
+        decomp_metas = [decomp for decomp in decomp_metas if decomp.is_complex]
+        if not decomp_metas:
+            return None, "No decompositions registered for the complex-valued input"
+    else:
+        decomp_metas = [decomp for decomp in decomp_metas if not decomp.is_complex]
+        if not decomp_metas:
+            return None, "No decompositions registered for the real-valued input"
+
+    # NOTE: Complex overload type matching logic has been removed to keep this simple
+    # There should no longer be overloads (for the same opset version) in torchlib anymore
+    return (decomp_metas[0].onnx_function, "The first implementation is used")

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_dynamic_shapes.py

@@ -0,0 +1,343 @@
+"""Compatibility functions for the torch.onnx.export API."""
+
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import inspect
+import warnings
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.export.dynamic_shapes import _DimHint, Dim
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+from torch.utils import _pytree
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+def from_dynamic_axes_to_dynamic_shapes(
+    model,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    *,
+    dynamic_axes=None,
+    output_names: set[str],
+    input_names: Sequence[str] | None = None,
+) -> tuple[dict[str, Any | None] | None, tuple[Any, ...], dict[str, Any] | None]:
+    """
+    Converts dynamic_axes into dynamic_shapes by wrapping the axis names with ``torch.export.Dim.DYNAMIC``.
+
+    dynamic_axes examples:
+    (1) dynamic_axes = {"x": {0: "my_custom_axis_name_1"}, "y": {1: "my_custom_axis_name_2"}}
+    (2) dynamic_axes = {"x": [0], "y": [1]}
+
+    these will be converted to dynamic_shapes respectively:
+    (1) dynamic_shapes = {"x": {0: Dim.DYNAMIC}, "y": {1: Dim.DYNAMIC}}
+    (2) dynamic_shapes = {"x": {0: Dim.DYNAMIC}, "y": {1: Dim.DYNAMIC}}
+
+    Detail on Dim.DYNAMIC: `#133620 <https://github.com/pytorch/pytorch/pull/133620>`_
+    """
+
+    warnings.warn(
+        "from_dynamic_axes_to_dynamic_shapes is deprecated and will be removed in a future release. "
+        "This function converts 'dynamic_axes' format (including custom axis names) to 'dynamic_shapes' format. "
+        "Instead of relying on this conversion, provide 'dynamic_shapes' directly with custom names.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+    # https://github.com/pytorch/pytorch/pull/128371
+    # 1. The function does not need to provide dynamic_shapes to torch.export.export
+    if dynamic_axes is None:
+        return None, args, kwargs
+
+    if input_names is None:
+        input_names = []
+
+    if kwargs is None:
+        kwargs = {}
+
+    dynamic_shapes: dict[str, Any | None] = {}
+    for input_name, axes in dynamic_axes.items():
+        # NOTE: torch.export.Dim.DYNAMIC does its best to infer the min and max values
+        # from the model, but it's not guaranteed to be dynamic.
+        if input_name in output_names:
+            # output names are not needed for dynamic_shapes
+            continue
+        if isinstance(axes, dict):
+            if any(not isinstance(k, int) for k in axes):
+                raise ValueError(
+                    "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]."
+                )
+            # str will be converted to Dim.DYNAMIC in convert_str_to_export_dim
+            dynamic_shapes[input_name] = axes
+        elif isinstance(axes, list):
+            if any(not isinstance(k, int) for k in axes):
+                raise ValueError(
+                    "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]."
+                )
+            dynamic_shapes[input_name] = dict.fromkeys(axes, torch.export.Dim.DYNAMIC)
+        elif axes is None:
+            dynamic_shapes[input_name] = None
+        else:
+            raise ValueError(
+                "Unsupported dynamic_axes format. Please provide a dict or a list."
+            )
+
+    for input_name in input_names:
+        if input_name not in dynamic_shapes:
+            dynamic_shapes[input_name] = None
+
+    # Order the inputs according to the signature of the model
+    sig = _signature(model)
+    inputs = []
+    for idx, param_name in enumerate(sig.parameters):
+        if idx < len(args):
+            inputs.append(args[idx])
+        elif param_name in kwargs:
+            inputs.append(kwargs[param_name])
+
+    # We need tree structure to represent dynamic_shapes
+    dynamic_shapes = _unflatten_dynamic_shapes_with_inputs_tree(inputs, dynamic_shapes)
+
+    # Since the dynamic_shapes are now in the order of the model parameters,
+    # we need to convert args and kwargs to the order of the model parameters.
+    return dynamic_shapes, tuple(inputs), {}
+
+
+def from_dynamic_shapes_to_dynamic_axes(
+    dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any],
+    input_names: Sequence[str],
+    exception: Exception,
+) -> dict[str, Any] | None:
+    """
+    Converts dynamic_shapes into dynamic_axes by removing torch.export.Dim wrapping
+    and converting to list or dict form based on whether dimension names are present.
+
+    dynamic_shapes examples:
+    (1) dynamic_shapes = {"x": {0: Dim("my_custom_axis_name_1")}, "y": {1: Dim("my_custom_axis_name_2")}}
+    (2) dynamic_shapes = ({0: Dim("my_custom_axis_name_1"}, {1: Dim("my_custom_axis_name_2")})
+
+    these will be converted to dynamic_axes respectively:
+    (1) dynamic_axes = {"x": [0], "y": [1]}
+    (2) dynamic_axes = {"x": [0], "y": [1]}
+
+    NOTE: If the model input is nested, so is the dynamic_shapes, we need to flatten the dynamic_shapes,
+    and then assign the axes to the input names in the order they are provided.
+
+    NOTE: input_names are used to assign the axes to the correct input names. If the input names are not
+    provided, or less than the dynamic inputs/axes, it raises an error.
+    """
+
+    flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes)
+
+    if len(input_names) < len(flat_dynamic_shapes):
+        raise ValueError(
+            "To construct dynamic_axes from dynamic_shapes, "
+            f"number of input names ({len(input_names)}) should be greater than or equal to "
+            f"the number of graph inputs(flat) ({len(flat_dynamic_shapes)})"
+        ) from exception
+
+    dynamic_axes: dict[str, list[int]] = {}
+    # input names are assigned in order
+    for input_name, axes in zip(input_names, flat_dynamic_shapes):
+        if axes is None:
+            continue
+
+        converted_axes: list[int] = []
+        if isinstance(axes, dict):
+            for axis, dim in axes.items():
+                if dim is None:
+                    continue
+                converted_axes.append(axis)
+            dynamic_axes[input_name] = converted_axes
+        elif isinstance(axes, (list, tuple)):
+            for idx, dim in enumerate(axes):
+                if dim is None:
+                    continue
+                converted_axes.append(idx)
+            dynamic_axes[input_name] = converted_axes
+    return dynamic_axes
+
+
+def _any_str_or_dim_in_dynamic_shapes(
+    dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any],
+) -> bool:
+    """Check if there is any string or Dim in the dynamic_shapes."""
+    flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes)
+    # This indicates the dynamic_shapes includes something we don't support in axes, and it's flattened
+    # to itself. Otherwise, flat_dynamic_shapes should be a list of dict/list/tuple (or None).
+    if any(
+        not isinstance(axes, (dict, list, tuple)) and axes is not None
+        for axes in flat_dynamic_shapes
+    ):
+        return False
+    # both str and Dim can provide custom names
+    for axes in flat_dynamic_shapes:
+        if isinstance(axes, dict):
+            for dim in axes.values():
+                if isinstance(dim, (str, Dim)):
+                    return True
+        elif isinstance(axes, (list, tuple)):
+            for dim in axes:
+                if isinstance(dim, (str, Dim)):
+                    return True
+    return False
+
+
+def convert_str_to_export_dim(
+    dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None,
+) -> tuple[dict[str, Any] | tuple[Any, ...] | list[Any] | None, bool]:
+    # 1. If there is no string in dynamic_shapes, we do not touch dynamic_shapes
+    if dynamic_shapes is None or not _any_str_or_dim_in_dynamic_shapes(dynamic_shapes):
+        return dynamic_shapes, False
+    # 2. Convert "name" to Dim.DYNAMIC with flattening and identify if there is any string
+    #    to be replaced with Dim.DYNAMIC, and then unflatten it back to the original structure.
+    #    for example: {"y": {0: "dim_0"}, "x": {1: "dim_1"}}
+    #    to {"y": {0: Dim.DYNAMIC}, "x": {1: Dim.DYNAMIC}}
+    dynamic_shapes_with_export_dim: list[
+        list[Dim | _DimHint | None] | dict[int, Dim | _DimHint | None] | None
+    ] = []
+    flat_dynamic_shapes, tree_structure = _flatten_dynamic_shapes_to_axes(
+        dynamic_shapes
+    )
+    for axes in flat_dynamic_shapes:
+        if axes is None:
+            dynamic_shapes_with_export_dim.append(None)
+        elif isinstance(axes, dict):
+            converted_axes_dict: dict[int, Dim | _DimHint | None] = {}
+            for axis, dim in axes.items():
+                if isinstance(dim, str):
+                    converted_axes_dict[axis] = torch.export.Dim.DYNAMIC
+                else:
+                    converted_axes_dict[axis] = dim
+            dynamic_shapes_with_export_dim.append(converted_axes_dict)
+        elif isinstance(axes, (list, tuple)):
+            converted_axes_list: list[Dim | _DimHint | None] = []
+            for dim in axes:
+                if isinstance(dim, str):
+                    converted_axes_list.append(torch.export.Dim.DYNAMIC)
+                else:
+                    converted_axes_list.append(dim)
+            dynamic_shapes_with_export_dim.append(converted_axes_list)
+
+    dynamic_shapes_with_export_dim = _pytree.tree_unflatten(
+        dynamic_shapes_with_export_dim, tree_structure
+    )
+    return (
+        dynamic_shapes_with_export_dim,
+        True,
+    )
+
+
+def create_rename_mapping(
+    inputs, dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any]
+) -> dict[str, str]:
+    """Create a mapping from old names to new names for dynamic axes."""
+
+    # NOTE: There's no need to handle cases where kwargs are out of order with the model signature,
+    # as torch.export.export supports dynamism only when kwargs and dynamic_shapes are provided in order.
+    # Reference: https://github.com/pytorch/pytorch/blob/49082f9dba3b79a344cb03652972ddbe7c3729cc/torch/export/_trace.py#L2034
+
+    flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes)
+    if len(inputs) != len(flat_dynamic_shapes):
+        warnings.warn(
+            "# ONNX model has different number of inputs than the flatten dynamic_shapes. "
+            "The dynamic axes will not be renamed.",
+            UserWarning,
+            stacklevel=3,
+        )
+        return {}
+    rename_mapping: dict[str, str] = {}
+    # NOTE: We assume that the flat_dynamic_shapes is in the same order as the inputs
+    # When the axis is static, or it connects to _DimHint in dynamic shapes, we skip renaming
+    for idx, axes in enumerate(flat_dynamic_shapes):
+        input = inputs[idx]
+        if isinstance(axes, dict):
+            for dim, axis in axes.items():
+                if not isinstance(input.shape[dim], ir.SymbolicDim):
+                    continue
+                old_name = input.shape[dim].value
+                if old_name is None:
+                    continue
+                # _DimHint, int and None exists in dynamic shapes, we skip renaming
+                if isinstance(axis, (_DimHint, int)) or axis is None:
+                    continue
+                # NOTE: ExportedProgram could give the axes the same name if they share
+                # the same shape constraints.
+                custom_name = _get_custom_axis_name(axis)
+                if input.shape[dim].value in rename_mapping:
+                    warnings.warn(
+                        f"# The axis name: {custom_name} will not be used, since it shares "
+                        f"the same shape constraints with another axis: {rename_mapping[input.shape[dim].value]}.",
+                        stacklevel=2,
+                    )
+                    continue
+                rename_mapping[input.shape[dim].value] = custom_name
+        elif isinstance(axes, (list, tuple)):
+            for dim, axis in enumerate(axes):
+                if not isinstance(input.shape[dim], ir.SymbolicDim):
+                    continue
+                old_name = input.shape[dim].value
+                if old_name is None:
+                    continue
+                # _DimHint, int and None exists in dynamic shapes, we skip renaming
+                if isinstance(axis, (_DimHint, int)) or axis is None:
+                    continue
+                # NOTE: ExportedProgram could give the axes the same name if they share
+                # the same shape constraints.
+                custom_name = _get_custom_axis_name(axis)
+                if input.shape[dim].value in rename_mapping:
+                    warnings.warn(
+                        f"# The axis name: {custom_name} will not be used, since it shares "
+                        f"the same shape constraints with another axis: {rename_mapping[input.shape[dim].value]}.",
+                        UserWarning,
+                        stacklevel=3,
+                    )
+                    continue
+                rename_mapping[input.shape[dim].value] = _get_custom_axis_name(axis)
+    return rename_mapping
+
+
+def _get_custom_axis_name(axis: Dim | str) -> str:
+    """Get the custom axis name from a torch.export.Dim."""
+    if isinstance(axis, Dim):
+        return axis.__name__
+    return axis
+
+
+def _unflatten_dynamic_shapes_with_inputs_tree(
+    inputs: list[Any],
+    dynamic_shapes: dict[str, Any],
+) -> dict[str, Any | None]:
+    _, tree_structure = _pytree.tree_flatten(inputs)
+    return _pytree.tree_unflatten(dynamic_shapes.values(), tree_structure)
+
+
+def _flatten_dynamic_shapes_to_axes(
+    dynamic_shapes: dict[str, Any | None] | tuple[Any, ...] | list[Any],
+) -> tuple[list[Any], _pytree.TreeSpec]:
+    # If it's a dict/list/tuple with torch.export.Dim, we consider it's an axis to dim mapping
+    def is_axes(x) -> bool:
+        return (
+            isinstance(x, dict)
+            and all(
+                isinstance(k, int)
+                and (v is None or isinstance(v, (Dim, _DimHint, str, int)))
+                for k, v in x.items()
+            )
+        ) or (
+            isinstance(x, (list, tuple))
+            and all(v is None or isinstance(v, (Dim, _DimHint, str, int)) for v in x)
+        )
+
+    return _pytree.tree_flatten(dynamic_shapes, is_leaf=is_axes)
+
+
+def _signature(model) -> inspect.Signature:
+    should_be_callable = getattr(model, "forward", model)
+    if callable(should_be_callable):
+        return inspect.signature(should_be_callable)
+    raise ValueError("model has no forward method and is not callable")

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_errors.py

@@ -0,0 +1,21 @@
+"""Error classes for the ONNX exporter."""
+
+from __future__ import annotations
+
+import torch.onnx.errors
+
+
+class TorchExportError(torch.onnx.errors.OnnxExporterError):
+    """Error during graph capturing using torch.export."""
+
+
+class ConversionError(torch.onnx.errors.OnnxExporterError):
+    """Error during ExportedProgram to ONNX conversion."""
+
+
+class DispatchError(ConversionError):
+    """Error during ONNX Function dispatching."""
+
+
+class GraphConstructionError(ConversionError):
+    """Error during ONNX graph construction."""

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_exportable_module.py

@@ -0,0 +1,196 @@
+"""Abstract interface for ONNX exportable modules."""
+
+from __future__ import annotations
+
+import abc
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class ExportableModule(torch.nn.Module, abc.ABC):
+    """Abstract interface for ONNX exportable modules.
+
+    Inherit from this class and implement the defined abstract methods
+    to create a module that can be exported to ONNX format.
+
+    Example::
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return x * 2
+
+
+        class MyExportableModule(torch.onnx.ExportableModule):
+            def __init__(self):
+                super().__init__()
+                self.model = Model()
+
+            def forward(self, x):
+                return self.model(x)
+
+            def example_arguments(self):
+                return (torch.randn(2, 3, 224, 224),), None
+
+            def input_names(self):
+                return ("input",)
+
+            def output_names(self):
+                return ("output",)
+
+            def dynamic_shapes(self):
+                return ({0: "batch_size"},)
+
+
+        exportable_module = MyExportableModule()
+        onnx_program = exportable_module.to_onnx()
+        # The model can also be supplied directly to torch.onnx.export
+        onnx_program = torch.onnx.export(exportable_module)
+    """
+
+    @abc.abstractmethod
+    def example_arguments(self) -> tuple[tuple[Any], dict[str, Any] | None]:
+        """Return example arguments for the model's forward method.
+
+        This method must be implemented by subclasses to provide sample inputs
+        that can be used for tracing, testing, and ONNX export. The returned
+        arguments should be representative of the expected input shapes and types
+        during inference.
+
+        Example::
+
+            def example_arguments(self):
+                # For a model expecting a single tensor input
+                return (torch.randn(1, 3, 224, 224),), None
+
+
+            def example_arguments(self):
+                # For a model with multiple inputs and keyword arguments
+                return (torch.randn(1, 3, 224, 224), torch.randn(1, 512)), {
+                    "temperature": 1.0
+                }
+
+        Returns:
+            A tuple containing:
+
+            - A tuple of positional arguments to pass to the forward method
+            - A dictionary of keyword arguments (or None if no kwargs are needed)
+        """
+        raise NotImplementedError
+
+    def dynamic_shapes(self) -> Any:
+        """Return dynamic shape specifications for the model's inputs.
+
+        Override this method to specify which dimensions of the input tensors
+        should be treated as dynamic during ONNX export. This allows the exported
+        model to accept inputs with varying sizes along the specified dimensions.
+
+        Example::
+
+            def dynamic_shapes(self):
+                # Specify batch dimension as dynamic for input named 'x'
+                return {"x": {0: "batch_size"}}
+
+
+            def dynamic_shapes(self):
+                # Multiple dynamic dimensions
+                return {
+                    "input": {0: "batch_size", 2: "height", 3: "width"},
+                    "mask": {0: "batch_size"},
+                }
+
+        Note:
+            The default implementation returns None, indicating all dimensions are static.
+
+        Returns:
+            Dynamic shape specification compatible with ``torch.export.export``.
+            Return None if all input dimensions should be static. The format can be:
+
+            - A dictionary mapping input names to dimension specifications
+            - A tuple/list of dimension specifications corresponding to inputs
+            - Any format accepted by the ``dynamic_shapes`` parameter of ``torch.export.export``
+        """
+        return None
+
+    def input_names(self) -> Sequence[str] | None:
+        """Return names for the model's input tensors.
+
+        Override this method to provide custom names for the input tensors in the
+        exported ONNX model. These names will be used as identifiers in the ONNX
+        graph and can be useful for debugging and model inspection.
+
+        Example::
+
+            def input_names(self):
+                return ["image", "mask"]
+
+
+            def input_names(self):
+                # For a single input
+                return ["input_tensor"]
+
+        Note:
+            The default implementation returns None, which results in auto-generated names.
+
+        Returns:
+            A sequence of strings representing input names, or None to use default names.
+            The number of names should match the number of positional arguments in the
+            forward method.
+        """
+        return None
+
+    def output_names(self) -> Sequence[str] | None:
+        """Return names for the model's output tensors.
+
+        Override this method to provide custom names for the output tensors in the
+        exported ONNX model. These names will be used as identifiers in the ONNX
+        graph and can be useful for debugging and model inspection.
+
+        Example::
+
+            def output_names(self):
+                return ["logits", "probabilities"]
+
+
+            def output_names(self):
+                # For a single output
+                return ["prediction"]
+
+        Note:
+            The default implementation returns None, which results in auto-generated names.
+
+        Returns:
+            A sequence of strings representing output names, or None to use default names.
+            The number of names should match the number of outputs from the forward method.
+            For models returning multiple outputs, provide a name for each output.
+        """
+        return None
+
+    def to_onnx(self, **kwargs: Any) -> torch.onnx.ONNXProgram:
+        """Export the module to ONNX format.
+
+        This method provides a convenient wrapper around ``torch.onnx.export`` that
+        automatically uses the example arguments, dynamic shapes, and input/output
+        names defined by the module. Additional export options can be specified via
+        keyword arguments.
+
+        See Also: ``torch.onnx.export`` for complete documentation of export options.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to ``torch.onnx.export``.
+                Common options include:
+
+                - ``opset_version`` (int): The ONNX opset version to target
+                - ``optimize`` (bool): Whether to apply optimizations to the exported model
+
+        Returns:
+            An ONNXProgram object containing the exported model and metadata.
+        """
+        result = torch.onnx.export(self, **kwargs)
+        if result is None:
+            raise AssertionError("result must be non-None")
+        return result

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_flags.py

@@ -0,0 +1,32 @@
+"""Internal flags for ONNX export."""
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, TypeVar
+from typing_extensions import ParamSpec
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+_is_onnx_exporting = False
+
+# Use ParamSpec to preserve parameter types instead of erasing to Any
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def set_onnx_exporting_flag(func: Callable[_P, _R]) -> Callable[_P, _R]:
+    @functools.wraps(func)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        global _is_onnx_exporting
+        _is_onnx_exporting = True
+        try:
+            return func(*args, **kwargs)
+        finally:
+            # Ensure it resets even if an exception occurs
+            _is_onnx_exporting = False
+
+    return wrapper

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_fx_passes.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import torch
+import torch.export
+import torch.fx
+from torch.onnx._internal.exporter import _decomp, _registration
+from torch.onnx._internal.fx import passes
+
+
+def decompose_with_registry(
+    exported_program: torch.export.ExportedProgram, registry: _registration.ONNXRegistry
+) -> torch.export.ExportedProgram:
+    """Decompose the exported program with the given registry.
+
+    This function is needed so it shows clearly on the profiler results.
+    """
+    onnx_registered_ops = set(_decomp.get_onnx_implemented_overloads(registry))
+    decomp_table = _decomp.create_onnx_friendly_decomposition_table(onnx_registered_ops)
+    return exported_program.run_decompositions(decomp_table)
+
+
+def insert_type_promotion_nodes(
+    graph_module: torch.fx.GraphModule,
+) -> None:
+    """Inplace pass to insert explicit type promotion nodes, recursively through nested modules."""
+    for module in graph_module.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            raise AssertionError(f"Expected GraphModule, got {type(module)}")
+        passes.InsertTypePromotion(module).run()
+
+
+def remove_assertion_nodes(graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    """Remove all assertion and check nodes from the FX graph"""
+    aten_assertion_targets = {
+        torch.ops.aten.sym_constrain_range_for_size.default,
+        torch.ops.aten._assert_async.default,
+        torch.ops.aten._assert_async.msg,
+        torch.ops.aten._assert_scalar.default,
+        torch.ops.aten._assert_tensor_metadata.default,
+    }
+    for gm in graph_module.modules():
+        for node in gm.graph.nodes:  # type: ignore[union-attr]
+            if node.op == "call_function" and node.target in aten_assertion_targets:
+                gm.graph.erase_node(node)  # type: ignore[operator, union-attr]
+        gm.recompile()  # type: ignore[operator]
+    return graph_module

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_ir_passes.py

@@ -0,0 +1,166 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING
+
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+from torch.onnx._internal.exporter import _constants
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+# The opset domain for ONNX operators
+_ONNX_DOMAIN = ""
+
+logger = logging.getLogger(__name__)
+
+
+def rename_inputs(model: ir.Model, new_names: Sequence[str]) -> None:
+    unique_names = frozenset(new_names)
+    if len(unique_names) != len(new_names):
+        seen = set()
+        duplicates = []
+        for name in new_names:
+            if name in seen:
+                duplicates.append(name)
+            seen.add(name)
+        raise ValueError(f"Input names cannot be duplicated: {duplicates}")
+
+    for input, new_name in zip(model.graph.inputs, new_names):
+        input.metadata_props["pkg.torch.onnx.original_node_name"] = str(input.name)
+        input.name = new_name
+
+
+def rename_outputs(model: ir.Model, new_names: Sequence[str]) -> None:
+    unique_names = frozenset(new_names)
+    if len(unique_names) != len(new_names):
+        seen = set()
+        duplicates = []
+        for name in new_names:
+            if name in seen:
+                duplicates.append(name)
+            seen.add(name)
+        raise ValueError(f"Output names cannot be duplicated: {duplicates}")
+
+    for output, new_name in zip(model.graph.outputs, new_names):
+        output.metadata_props["pkg.torch.onnx.original_node_name"] = str(output.name)
+        output.name = new_name
+
+
+def _all_values(model: ir.Model):
+    """Yield all values in a model."""
+    # Yield all values in the model
+    yield from model.graph.inputs
+    yield from model.graph.initializers.values()
+    for node in ir.traversal.RecursiveGraphIterator(model.graph):
+        yield from node.outputs
+    # Yield all values in functions
+    for function in model.functions.values():
+        yield from function.inputs
+        for node in ir.traversal.RecursiveGraphIterator(function):
+            yield from node.outputs
+
+
+def _replace_names(shape_expr: str, rename_mapping: dict[str, str]) -> str:
+    """Replace all known names in a shape expression with new names."""
+    for old_name, new_name in rename_mapping.items():
+        shape_expr = re.sub(
+            rf"(?<!\w){re.escape(old_name)}(?!\w)", new_name, shape_expr
+        )
+    return shape_expr
+
+
+def rename_axis(
+    model: ir.Model, rename_mapping: dict[str | ir.SymbolicDim, str]
+) -> None:
+    """Rename dynamic axes in a model according to the specified dynamic_axes names."""
+
+    # Create a mapping from string to string for easier replacement
+    string_mapping: dict[str, str] = {}
+    for key, value in tuple(rename_mapping.items()):
+        if isinstance(key, ir.SymbolicDim):
+            if isinstance(key.value, str):
+                string_mapping[key.value] = value
+            else:
+                raise ValueError(
+                    f"Invalid SymbolicDim value in rename_mapping: {key.value!r}. "
+                    "Expected str."
+                )
+        elif isinstance(key, str):
+            string_mapping[key] = value
+        else:
+            raise ValueError(
+                f"Invalid key type in rename_mapping: {type(key)}({key!r}). Expected "
+                "str or ir.SymbolicDim."
+            )
+
+    # NOTE: Mapping needs to be sorted by length because the shape expression
+    # could have multiple ways to be expressed, for example,
+    # {"s1": sequence_length, "s11": "past_sequence_length", "s1 + s11": "masked_sequence_length"}
+    # We prefer the replacement starts from the longest match.
+    sorted_rename_mapping = dict(
+        sorted(string_mapping.items(), key=lambda item: len(item[0]), reverse=True)
+    )
+    for value in _all_values(model):
+        if value.shape is None:
+            continue
+        new_shape = []
+        changed = False
+        for dim in value.shape:
+            if not isinstance(dim, ir.SymbolicDim):
+                new_shape.append(dim)
+                continue
+            dim_name = dim.value
+            if dim_name in sorted_rename_mapping:
+                new_shape.append(sorted_rename_mapping[dim_name])
+                changed = True
+            elif dim_name is not None:
+                # For example: "2*s1", "s1+1", "s1-1", "s1*s2", "s1/s2"
+                new_name = _replace_names(dim_name, sorted_rename_mapping)
+                new_shape.append(new_name)
+                if new_name != dim_name:
+                    changed = True
+            else:
+                new_shape.append(None)
+        if changed:
+            value.shape = ir.Shape(new_shape)
+
+
+def _maybe_set_opset_version(
+    opset_imports: dict[str, int], domain: str, version: int | None
+) -> None:
+    """Set the opset version for the domain."""
+    if domain in opset_imports and opset_imports[domain] != 1:
+        # Already set
+        return
+    if domain == _ONNX_DOMAIN:
+        opset_imports[domain] = _constants.TORCHLIB_OPSET
+        return
+    if version is None:
+        # We don't know the opset version, so set it to 1
+        # This is valid for the custom function domains like "pkg.torch.__subgraph__"
+        opset_imports[domain] = 1
+        return
+    # Set the known opset version for the domain
+    opset_imports[domain] = version
+
+
+def add_opset_imports(model: ir.Model) -> None:
+    """Collect all opsets used and add opset imports to the model and functions."""
+    for node in ir.traversal.RecursiveGraphIterator(model.graph):
+        domain = node.domain
+        _maybe_set_opset_version(model.opset_imports, domain, node.version)
+
+    for function in model.functions.values():
+        for node in ir.traversal.RecursiveGraphIterator(function):
+            domain = node.domain
+            _maybe_set_opset_version(function.opset_imports, domain, node.version)
+        for domain, version in function.opset_imports.items():
+            # Add all opsets used in the function to the model, because ONNX Runtime
+            # does not handle adding the opset imports to the model after inlining during inference.
+            # This should happen after all opsets are collected for the function from its nodes.
+            _maybe_set_opset_version(model.opset_imports, domain, version)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_isolated.py

@@ -0,0 +1,65 @@
+"""Isolated calls to methods that may segfault."""
+
+from __future__ import annotations
+
+import multiprocessing
+import os
+import warnings
+from typing import Any, TYPE_CHECKING, TypeVar
+from typing_extensions import ParamSpec, TypeVarTuple, Unpack
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+_Ts = TypeVarTuple("_Ts")
+
+_IS_WINDOWS = os.name == "nt"
+
+
+def _call_function_and_return_exception(
+    func: Callable[[Unpack[_Ts]], _R], args: tuple[Unpack[_Ts]], kwargs: dict[str, Any]
+) -> _R | Exception:
+    """Call function and return a exception if there is one."""
+
+    try:
+        return func(*args, **kwargs)
+    except Exception as e:
+        return e
+
+
+def safe_call(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
+    """Call a function in a separate process.
+
+    Args:
+        func: The function to call.
+        args: The positional arguments to pass to the function.
+        kwargs: The keyword arguments to pass to the function.
+
+    Returns:
+        The return value of the function.
+
+    Raises:
+        Exception: If the function raised an exception.
+    """
+    if _IS_WINDOWS:
+        # On Windows, we cannot create a new process with fork.
+        warnings.warn(
+            f"A new process is not created for {func} on Windows.", stacklevel=1
+        )
+        return func(*args, **kwargs)
+
+    with multiprocessing.get_context("fork").Pool(1) as pool:
+        # It is important to fork a process here to prevent the main logic from
+        # running again when the user does not place it under a `if __name__ == "__main__":`
+        # block.
+        result = pool.apply_async(
+            _call_function_and_return_exception, (func, args, kwargs)
+        )
+        result = result.get(timeout=5)
+    if isinstance(result, Exception):
+        raise result
+    return result

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_onnx_program.py

@@ -0,0 +1,521 @@
+# mypy: allow-untyped-defs
+# mypy: disable-error-code="attr-defined,name-defined"
+from __future__ import annotations
+
+
+__all__ = ["ONNXProgram"]
+
+import contextlib
+import copy
+import gc
+import logging
+import os
+import tempfile
+import textwrap
+import warnings
+from collections.abc import Callable, Sequence
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.onnx._internal._lazy_import import onnx, onnx_ir as ir, onnxscript_apis
+from torch.onnx._internal.exporter import _dynamic_shapes, _ir_passes
+from torch.utils import _pytree
+
+
+# NOTE: DO NOT import module from torch.onnx._internal to this module in the global scope
+# because ONNXProgram is exposed to the public API
+
+if TYPE_CHECKING:
+    import numpy as np
+    import onnxruntime as ort
+
+_LARGE_MODEL_THRESHOLD = 1536 * 1024 * 1024  # 1536MB
+_NP_UNSUPPORTED_DTYPES_8BIT = frozenset(
+    {
+        torch.float8_e4m3fn,
+        torch.float8_e4m3fnuz,
+        torch.float8_e5m2,
+        torch.float8_e5m2fnuz,
+    }
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _ort_session_initializer(model: str | bytes) -> ort.InferenceSession:
+    """Initialize an ONNX Runtime inference session with the specified model."""
+    import onnxruntime as ort
+
+    session_options = ort.SessionOptions()
+    session_options.log_severity_level = 3  # 3: Error
+    possible_providers = (
+        "CUDAExecutionProvider",
+        "CPUExecutionProvider",
+    )
+    available_providers = set(ort.get_available_providers())
+    providers = [
+        provider for provider in possible_providers if provider in available_providers
+    ]
+    return ort.InferenceSession(
+        model, providers=providers, sess_options=session_options
+    )
+
+
+def _count_initializer_size(graph: ir.Graph) -> int:
+    """Count the total size of the initializers in bytes."""
+    return sum(
+        v.const_value.nbytes
+        for v in graph.initializers.values()
+        if v.const_value is not None
+    )
+
+
+@contextlib.contextmanager
+def _set_graph_outputs(
+    graph: ir.Graph,
+    outputs: list[ir.Value],
+):
+    """Temporarily set the outputs of the graph.
+
+    Args:
+        graph: The graph to set the outputs for.
+        outputs: The outputs to set.
+    """
+    original_outputs = list(graph.outputs)
+    graph.outputs.clear()
+    graph.outputs.extend(outputs)
+    try:
+        yield
+    finally:
+        graph.outputs.clear()
+        graph.outputs.extend(original_outputs)
+
+
+def _create_value_mapping(graph: ir.Graph) -> dict[str, ir.Value]:
+    """Return a dictionary mapping names to values in the graph.
+
+    The mapping does not include values from subgraphs.
+
+    Args:
+        graph: The graph to extract the mapping from.
+
+    Returns:
+        A dictionary mapping names to values.
+    """
+    values: dict[str, ir.Value] = {}
+    values.update(graph.initializers)
+    # The names of the values can be None or "", which we need to exclude
+    for input in graph.inputs:
+        if not input.name:
+            continue
+        values[input.name] = input
+    for node in graph:
+        for value in node.outputs:
+            if not value.name:
+                continue
+            values[value.name] = value
+    return values
+
+
+def _to_numpy_array(input: torch.Tensor | int | float | str | bool) -> np.ndarray:
+    if isinstance(input, (int, float, str, bool)):
+        return ir.tensor(input).numpy()
+
+    from torch.onnx._internal.exporter import _core
+
+    return _core.TorchTensor(input).numpy()
+
+
+def _from_numpy_array(array: np.ndarray) -> torch.Tensor:
+    """Convert a NumPy array to a PyTorch tensor."""
+    import ml_dtypes  # type: ignore[import-not-found]
+    import numpy as np
+
+    if array.dtype == ml_dtypes.bfloat16:
+        return torch.from_numpy(array.view(np.uint16)).view(torch.bfloat16)
+    if array.dtype == ml_dtypes.float8_e4m3fn:
+        return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e4m3fn)
+    if array.dtype == ml_dtypes.float8_e4m3fnuz:
+        return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e4m3fnuz)
+    if array.dtype == ml_dtypes.float8_e5m2:
+        return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e5m2)
+    if array.dtype == ml_dtypes.float8_e5m2fnuz:
+        return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e5m2fnuz)
+    return torch.from_numpy(array)
+
+
+def _to_ort_value(input: torch.Tensor | int | float | str | bool) -> ort.OrtValue:
+    """Convert a PyTorch tensor to an ONNX Runtime OrtValue."""
+    import numpy as np
+    import onnxruntime as ort
+
+    from torch.onnx._internal.exporter import _core
+
+    if isinstance(input, (int, float, str, bool)):
+        # Convert scalar values to OrtValue
+        dtype_mapping = {
+            int: np.int64,
+            float: np.float32,
+        }
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
+        dtype = dtype_mapping.get(type(input))
+        return ort.OrtValue.ortvalue_from_numpy(np.array(input, dtype=dtype))
+
+    if input.dtype == torch.bfloat16 or input.dtype in _NP_UNSUPPORTED_DTYPES_8BIT:
+        if hasattr(ort.OrtValue, "ortvalue_from_numpy_with_onnx_type"):
+            # This requires ONNX Runtime 1.21 or newer
+            if input.dtype == torch.bfloat16:
+                uint_type = torch.uint16
+            else:
+                uint_type = torch.uint8
+            onnx_type = _core.torch_dtype_to_onnx_dtype(input.dtype)
+            # Make tensor contiguous to ensure view() works
+            input = input.contiguous()
+            return ort.OrtValue.ortvalue_from_numpy_with_onnx_type(
+                input.view(uint_type).numpy(force=True), onnx_element_type=onnx_type
+            )
+        raise RuntimeError(
+            f"Failed to convert tensor of type '{input.dtype}' to OrtValue. "
+            "Please ensure that ONNX Runtime is built with DLPack support or is the latest version"
+        )
+    # TODO(#151064): Use dlpack when ORT properly supports it
+    return ort.OrtValue.ortvalue_from_numpy(input.numpy(force=True))
+
+
+def _from_ort_value(value: ort.OrtValue) -> torch.Tensor:
+    if value.element_type() in (
+        ir.DataType.BFLOAT16,
+        ir.DataType.FLOAT8E4M3FN,
+        ir.DataType.FLOAT8E4M3FNUZ,
+        ir.DataType.FLOAT8E5M2,
+        ir.DataType.FLOAT8E5M2FNUZ,
+    ):
+        # This requires ONNX Runtime 1.21 or newer
+        try:
+            return torch.from_dlpack(value._get_c_value())
+        except Exception as e:
+            raise RuntimeError(
+                "Failed to convert OrtValue to torch.Tensor. "
+                "Please ensure that ONNX Runtime is built with DLPack support or is the latest version"
+            ) from e
+    return torch.from_numpy(value.numpy())
+
+
+class ONNXProgram:
+    """A class to represent an ONNX program that is callable with torch tensors.
+
+    Attributes:
+        model: The ONNX model as an ONNX IR model object.
+        exported_program: The exported program that produced the ONNX model.
+    """
+
+    def __init__(
+        self, model: ir.Model, exported_program: torch.export.ExportedProgram | None
+    ) -> None:
+        """Initialize the ONNX program with the specified model and exported program.
+        Args:
+            model: The ONNX model.
+            exported_program: The exported program that produced the ONNX model. Optional.
+        """
+        self.model: ir.Model = model
+        self.exported_program = exported_program
+        self._inference_session: ort.InferenceSession | None = None
+        self._tempdir: tempfile.TemporaryDirectory | None = None
+        # Strategy used to capture the exported program
+        self._capture_strategy: str | None = None
+
+    def __repr__(self) -> str:
+        return f"""\
+ONNXProgram(
+    model=
+{textwrap.indent(str(self.model), " " * 8)}
+    ,
+    exported_program=
+{textwrap.indent(str(self.exported_program), " " * 8)}
+)
+"""
+
+    def __call__(self, *args, **kwargs) -> Sequence[torch.Tensor]:
+        """Run the ONNX model with the same arguments you would provide to the GraphModule."""
+        import onnxruntime as ort
+
+        flatten_args = _process_args(args, kwargs)
+
+        if self._inference_session is None:
+            self.initialize_inference_session()
+
+        if self._inference_session is None:
+            raise AssertionError("_inference_session must be non-None")
+
+        ort_input = {
+            k.name: _to_ort_value(v)
+            for k, v in zip(self.model.graph.inputs, flatten_args)
+        }
+        run_options = ort.RunOptions()
+        run_options.log_severity_level = 3  # 3: Error
+        logger.debug("Running the inference session with %s arguments.", len(ort_input))
+        outputs = self._inference_session.run_with_ort_values(
+            None, ort_input, run_options=run_options
+        )
+        logger.debug("Inference session run completed.")
+        return tuple(_from_ort_value(output) for output in outputs)
+
+    def call_reference(self, *args, **kwargs) -> Sequence[torch.Tensor]:
+        """Run the ONNX model using the reference backend."""
+        import onnx.reference
+
+        evaluator = onnx.reference.ReferenceEvaluator(self.model_proto)
+
+        flatten_args = _process_args(args, kwargs)
+        ref_input = {
+            k.name: _to_numpy_array(v)
+            for k, v in zip(self.model.graph.inputs, flatten_args)
+        }
+        outputs = evaluator.run(None, ref_input)  # type: ignore[arg-type]
+        if not isinstance(outputs, Sequence):
+            raise AssertionError(f"Expected Sequence, got {type(outputs)}")
+        return tuple(_from_numpy_array(output) for output in outputs)
+
+    def compute_values(
+        self, value_names: Sequence[str], args=(), kwargs=None
+    ) -> Sequence[torch.Tensor]:
+        """Compute the values of the specified names in the ONNX model.
+
+        This method is used to compute the values of the specified names in the ONNX model.
+        The values are returned as a dictionary mapping names to tensors.
+
+        Args:
+            value_names: The names of the values to compute.
+
+        Returns:
+            A dictionary mapping names to tensors.
+        """
+        if kwargs is None:
+            kwargs = {}
+        self.release()
+        values = _create_value_mapping(self.model.graph)
+        for name in value_names:
+            if name not in values:
+                raise ValueError(
+                    f"Value '{name}' not found in the model. "
+                    "Please provide a valid value name."
+                )
+        temporary_outputs = [values[name] for name in value_names]
+        with _set_graph_outputs(self.model.graph, temporary_outputs):
+            try:
+                result = self(*args, **kwargs)
+            finally:
+                self.release()
+        return result
+
+    @property
+    def model_proto(self) -> onnx.ModelProto:
+        """Return the ONNX ``ModelProto`` object."""
+        return ir.serde.serialize_model(self.model)
+
+    def optimize(self) -> None:
+        """Optimize the ONNX model.
+
+        This method optimizes the ONNX model by performing constant folding and
+        eliminating redundancies in the graph. The optimization is done in-place.
+        """
+        self.model = onnxscript_apis.optimize(self.model)
+
+    def save(
+        self,
+        destination: str | os.PathLike,
+        *,
+        include_initializers: bool = True,
+        keep_initializers_as_inputs: bool = False,
+        external_data: bool | None = None,
+    ) -> None:
+        """Save the ONNX model to the specified destination.
+
+        When ``external_data`` is ``True`` or the model is larger than 2GB,
+        the weights are saved as external data in a separate file.
+
+        Initializer (model weights) serialization behaviors:
+
+        * ``include_initializers=True``, ``keep_initializers_as_inputs=False`` (default):
+          The initializers are included in the saved model.
+        * ``include_initializers=True``, ``keep_initializers_as_inputs=True``:
+          The initializers are included in the saved model and kept as model inputs.
+          Choose this option if you want the ability to override the model weights
+          during inference.
+        * ``include_initializers=False``, ``keep_initializers_as_inputs=False``:
+          The initializers are not included in the saved model and are not listed
+          as model inputs. Choose this option if you want to attach the initializers
+          to the ONNX model in a separate, post-processing, step.
+        * ``include_initializers=False``, ``keep_initializers_as_inputs=True``:
+          The initializers are not included in the saved model but are listed as model
+          inputs. Choose this option if you want to supply the initializers during
+          inference and want to minimize the size of the saved model.
+
+        Args:
+            destination: The path to save the ONNX model to.
+            include_initializers: Whether to include the initializers in the saved model.
+            keep_initializers_as_inputs: Whether to keep the initializers as inputs in the saved model.
+                If `True`, the initializers are added as inputs to the model which means they can be overwritten.
+                by providing the initializers as model inputs.
+            external_data: Whether to save the weights as external data in a separate file.
+
+        Raises:
+            TypeError: If ``external_data`` is ``True`` and ``destination`` is not a file path.
+        """
+        original_initializers = copy.copy(self.model.graph.initializers)
+        original_inputs = copy.copy(self.model.graph.inputs)
+
+        # Adjust the model based on options
+        if not include_initializers:
+            self.model.graph.initializers.clear()
+        if keep_initializers_as_inputs:
+            self.model.graph.inputs.extend(original_initializers.values())  # type: ignore[arg-type]
+
+        try:
+            # Save the model to disk
+            if (
+                external_data
+                or _count_initializer_size(self.model.graph) > _LARGE_MODEL_THRESHOLD
+            ):
+                onnxscript_apis.save_model_with_external_data(self.model, destination)
+            else:
+                ir.save(self.model, destination)
+        finally:
+            # Revert the changes to the model
+            if not include_initializers:
+                self.model.graph.initializers.update(original_initializers)
+            if keep_initializers_as_inputs:
+                self.model.graph.inputs.clear()
+                self.model.graph.inputs.extend(original_inputs)
+
+    def apply_weights(self, state_dict: dict[str, torch.Tensor]) -> None:
+        """Apply the weights from the specified state dict to the ONNX model.
+
+        Use this method to replace FakeTensors or other weights.
+
+        Args:
+            state_dict: The state dict containing the weights to apply to the ONNX model.
+        """
+        from torch.onnx._internal.exporter import _core
+
+        for name, tensor in state_dict.items():
+            if name in self.model.graph.initializers:
+                self.model.graph.initializers[name].const_value = _core.TorchTensor(
+                    tensor, name
+                )
+            else:
+                warnings.warn(
+                    f"Weight '{name}' not found in the model. Skipped applying.",
+                    category=torch.onnx.errors.OnnxExporterWarning,
+                    stacklevel=1,
+                )
+
+    def initialize_inference_session(
+        self,
+        initializer: Callable[
+            [str | bytes], ort.InferenceSession
+        ] = _ort_session_initializer,
+    ) -> None:
+        """Initialize the ONNX Runtime inference session.
+
+        Args:
+            initializer: The function to initialize the ONNX Runtime inference
+                session with the specified model. By default, it uses the
+                :func:`_ort_session_initializer` function.
+        """
+        # TODO(justinchuby): Allow different inference options
+        logger.debug("Initializing the inference session.")
+        if (
+            byte_size := _count_initializer_size(self.model.graph)
+        ) > _LARGE_MODEL_THRESHOLD:
+            logger.debug("The model initializers is larger than 1.5GB (%s).", byte_size)
+            # Save the model to a temporary file if too large
+            self._tempdir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
+            model_path = os.path.join(self._tempdir.name, "model.onnx")
+            self.save(model_path, external_data=True)
+            model = model_path
+        else:
+            model = self.model_proto.SerializeToString()  # type: ignore[assignment]
+
+        self._inference_session = initializer(model)
+        logger.debug("Inference session initialized.")
+
+    def release(self) -> None:
+        """Release the inference session.
+
+        You may call this method to release the resources used by the inference session.
+        """
+        # Release the inference session first so that the model file can be deleted
+        if self._inference_session is not None:
+            self._inference_session = None
+        gc.collect()
+        if self._tempdir is not None:
+            self._tempdir.cleanup()
+            self._tempdir = None
+
+    def rename_axes(self, rename_mapping: dict[str | ir.SymbolicDim, str]) -> None:
+        """Rename axes in a model according to the specified rename mapping.
+
+        Example::
+
+            batch = onnx_program.model.graph.inputs[0].shape[0]
+            seq_len = onnx_program.model.graph.inputs[0].shape[2]
+            rename_mapping = {
+                batch: "batch",
+                seq_len: "seq_len",
+            }
+            onnx_program.rename_axes(rename_mapping)
+
+        Args:
+            rename_mapping: A dictionary mapping old axes to new axis names.
+                Keys can be either:
+
+                * String axis names (e.g., "s1", "s2")
+                * SymbolicDim objects obtained from the model
+                  (e.g., onnx_program.model.graph.inputs[0].shape[0])
+
+                Values must be strings representing the new axis names.
+        """
+        _ir_passes.rename_axis(self.model, rename_mapping)
+
+    def _rename_dynamic_axes(
+        self,
+        dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any],
+    ) -> None:
+        """Rename dynamic axes in a model according to the specified dynamic_axes names."""
+        rename_mapping = _dynamic_shapes.create_rename_mapping(
+            self.model.graph.inputs, dynamic_shapes
+        )
+        _ir_passes.rename_axis(self.model, rename_mapping)
+
+
+def _process_args(args, kwargs) -> tuple[torch.Tensor, ...]:
+    """Process input arguments for the ONNX model."""
+    args = _flatten_inputs(args, kwargs)
+    args = _remove_none_from_inputs(args)
+    args = _convert_complex_to_real_representation(args)
+    return args
+
+
+def _flatten_inputs(model_args, model_kwargs):
+    flattened_args, _ = _pytree.tree_flatten((model_args, model_kwargs))
+    return flattened_args
+
+
+def _remove_none_from_inputs(model_args):
+    return tuple(arg for arg in model_args if arg is not None)
+
+
+def _convert_complex_to_real_representation(model_args):
+    """Convert complex dtype tensors to real representation tensors.
+
+    ONNX does not support complex dtype tensors. Thus, we convert complex dtype tensors
+    to real representation tensors (i.e., float dtype tensors with an extra dimension
+    representing the real and imaginary parts of the complex number).
+    """
+    return tuple(
+        torch.view_as_real(arg.resolve_conj())
+        if isinstance(arg, torch.Tensor) and arg.is_complex()
+        else arg
+        for arg in model_args
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_registration.py

@@ -0,0 +1,303 @@
+"""Module for handling ATen to ONNX functions registration.
+
+https://github.com/pytorch/pytorch/blob/6aa5bb1a76dee8112f1a9e7c194c790b5cdc6462/torch/onnx/_internal/fx/registration.py
+"""
+
+# NOTE: Why do we need a different registry than the one in torchlib?
+# The registry in torchlib is used to register functions that are already implemented in
+# torchlib, and is designed to be a static singleton. It does not take into account custom ops or different
+# opsets etc. The registry implemented for the exporter is designed to be modifiable at
+# export time by users, and is designed with dispatching in mind.
+
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import dataclasses
+import importlib.util
+import logging
+import math
+import operator
+import types
+from collections.abc import Callable
+from typing import Literal, TypeAlias
+
+import torch
+import torch._ops
+from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript, onnxscript_apis
+from torch.onnx._internal.exporter import _constants, _schemas
+from torch.onnx._internal.exporter._torchlib import _torchlib_registry
+
+
+TorchOp: TypeAlias = torch._ops.OpOverload | types.BuiltinFunctionType | Callable
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class OnnxDecompMeta:
+    """A wrapper of onnx-script function with additional metadata.
+
+    onnx_function: The onnx-script function from torchlib.
+    fx_target: The PyTorch node callable target.
+    signature: The ONNX signature of the function. When None, the signature is inferred.
+    is_custom: Whether the function is a custom function.
+    is_complex: Whether the function is a function that handles complex valued inputs.
+    opset_introduced:
+        The ONNX opset version in which the function was introduced.
+        Its specifies the minimum ONNX opset version required to use the function.
+    device: The device the function is registered to. If None, it is registered to all devices.
+    skip_signature_inference: Whether to skip signature inference for the function.
+    """
+
+    onnx_function: Callable
+    fx_target: TorchOp
+    signature: ir.schemas.OpSignature | None
+    is_custom: bool = False
+    is_complex: bool = False
+    opset_introduced: int = 18
+    device: Literal["cuda", "cpu"] | str | None = None  # noqa: PYI051
+    skip_signature_inference: bool = False
+
+    def __post_init__(self) -> None:
+        if self.signature is None and not self.skip_signature_inference:
+            try:
+                if isinstance(self.onnx_function, onnxscript.OnnxFunction):
+                    signature = _schemas.op_signature_from_function(
+                        self.onnx_function,
+                        self.onnx_function.function_ir.domain,
+                        self.onnx_function.name,
+                        since_version=self.onnx_function.opset.version,
+                    )
+                else:
+                    signature = _schemas.op_signature_from_function(
+                        self.onnx_function, "__traced", self.onnx_function.__name__
+                    )
+            except Exception as e:
+                # Log an warning if the op is custom. Raise exception for builtin ops.
+                if not self.is_custom:
+                    raise
+                else:
+                    # When the function is targeting an HOP, for example, it will accept
+                    # functions as arguments and fail to generate an ONNX signature.
+                    # In this case we set signature to None and dispatch to this function always.
+                    logger.warning(  # noqa: G200
+                        "Failed to infer the signature for function '%s' because '%s'"
+                        "All nodes targeting `%s` will be dispatched to this function",
+                        self.onnx_function,
+                        e,
+                        self.fx_target,
+                    )
+            else:
+                self.signature = signature
+                self.onnx_function._pt_onnx_signature = signature  # type: ignore[attr-defined]
+
+
+def _get_overload(qualified_name: str) -> torch._ops.OpOverload | None:
+    """Obtain the torch op from <namespace>::<op_name>[.<overload>]"""
+    # TODO(justinchuby): Handle arbitrary custom ops
+    namespace, opname_overload = qualified_name.split("::")
+    op_name, *maybe_overload = opname_overload.split(".", 1)
+    if namespace == "_operator":
+        # Builtin functions
+        return getattr(operator, op_name)
+    if namespace == "math":
+        return getattr(math, op_name)
+    if namespace == "torchvision":
+        if importlib.util.find_spec("torchvision") is None:
+            logger.warning("torchvision is not installed. Skipping %s", qualified_name)
+            return None
+    try:
+        op_packet = getattr(getattr(torch.ops, namespace), op_name)
+        if maybe_overload:
+            overload = maybe_overload[0]
+        elif "default" in op_packet._overload_names or "" in op_packet._overload_names:
+            # Has a default overload
+            overload = "default"
+        else:
+            logger.warning(
+                "'%s' does not have a 'default' overload. This could be an error in specifying the op name. Ignoring.",
+                qualified_name,
+                stacklevel=1,
+            )
+            return None
+
+        return getattr(op_packet, overload)  # type: ignore[call-overload]
+    except AttributeError:
+        if qualified_name.endswith("getitem"):
+            # This is a special case where we registered the function incorrectly,
+            # but for BC reasons (pt<=2.4) we need to keep it.
+            return None
+        logger.info("'%s' is not found in this version of PyTorch.", qualified_name)
+        return None
+    except Exception:
+        logger.exception("Failed to find torch op '%s'", qualified_name)
+        return None
+
+
+class ONNXRegistry:
+    """Registry for ONNX functions.
+
+    The registry maintains a mapping from qualified names to symbolic functions under a
+    fixed opset version. It supports registering custom onnx-script functions and for
+    dispatcher to dispatch calls to the appropriate function.
+
+    """
+
+    def __init__(self) -> None:
+        """Initializes the registry"""
+        self._opset_version = _constants.TORCHLIB_OPSET
+        self.functions: dict[TorchOp | str, list[OnnxDecompMeta]] = {}
+
+    @property
+    def opset_version(self) -> int:
+        """The ONNX opset version the exporter should target."""
+        return self._opset_version
+
+    @classmethod
+    def from_torchlib(cls, opset_version=_constants.TORCHLIB_OPSET) -> ONNXRegistry:
+        """Populates the registry with ATen functions from torchlib.
+
+        Args:
+            torchlib_registry: The torchlib registry to use for populating the registry.
+        """
+        registry = cls()
+        registry._opset_version = opset_version
+        for meta in _torchlib_registry.get_torchlib_ops():
+            registry._register(meta.fx_target, meta)
+
+        # TODO(justinchuby): Remove this once torchlib is migrated to PyTorch
+        torchlib_ops = onnxscript_apis.get_torchlib_ops()
+
+        for torchlib_meta in torchlib_ops:
+            qualified_name = torchlib_meta.qualified_name
+            overload_func = torchlib_meta.function
+            try:
+                # NOTE: This is heavily guarded with try-except because we don't want
+                # to fail the entire registry population if one function fails.
+                target = _get_overload(qualified_name)
+                if target is None:
+                    continue
+
+                meta = OnnxDecompMeta(
+                    onnx_function=overload_func,
+                    fx_target=target,
+                    signature=None,
+                    is_custom=False,
+                    is_complex=torchlib_meta.is_complex,
+                )
+                registry._register(target, meta)
+            except Exception:
+                logger.exception("Failed to register '%s'. Skipped", qualified_name)
+                continue
+
+        registry._cleanup_registry_based_on_opset_version()
+        return registry
+
+    def _register(
+        self,
+        target: TorchOp,
+        onnx_decomposition: OnnxDecompMeta,
+    ) -> None:
+        """Registers a OnnxDecompMeta to an operator.
+
+        Args:
+            target: The PyTorch node callable target.
+            onnx_decomposition: The OnnxDecompMeta to register.
+        """
+        target_or_name: str | TorchOp
+        if isinstance(target, torch._ops.OpOverload):
+            # Get the qualified name of the aten op because torch._ops.OpOverload lookup in
+            # a dictionary is unreliable for some reason.
+            target_or_name = target.name()
+        else:
+            target_or_name = target
+        if onnx_decomposition.is_custom:
+            self.functions.setdefault(target_or_name, []).insert(0, onnx_decomposition)
+        else:
+            self.functions.setdefault(target_or_name, []).append(onnx_decomposition)
+
+    def register_op(
+        self,
+        target: TorchOp,
+        function: Callable,
+        is_complex: bool = False,
+    ) -> None:
+        """Registers a custom operator: torch.ops.<namespace>.<op_name>.<overload>.
+
+        Args:
+            target: The PyTorch node callable target.
+            function: The onnx-script function to register.
+            is_complex: Whether the function is a function that handles complex valued inputs.
+        """
+        if isinstance(target, torch._ops.OpOverloadPacket):
+            raise TypeError(
+                f"Target '{target}' should be provided as an OpOverload instead of an "
+                "OpOverloadPacket. You can get the default overload with "
+                "<op>.default"
+            )
+
+        self._register(
+            target,
+            OnnxDecompMeta(
+                onnx_function=function,
+                fx_target=target,
+                signature=None,
+                is_custom=True,
+                is_complex=is_complex,
+            ),
+        )
+
+    def get_decomps(self, target: TorchOp) -> list[OnnxDecompMeta]:
+        """Returns a list of OnnxDecompMeta for the given op: torch.ops.<namespace>.<op_name>.<overload>.
+
+        The list is ordered by the time of registration. The custom operators should come
+        first in the list.
+
+        Args:
+            target: The PyTorch node callable target.
+        Returns:
+            A list of OnnxDecompMeta corresponding to the given name, or None if
+            the name is not in the registry.
+        """
+        target_or_name: str | TorchOp
+        if isinstance(target, torch._ops.OpOverload):
+            # Get the qualified name of the aten op because torch._ops.OpOverload lookup in
+            # a dictionary is unreliable for some reason.
+            target_or_name = target.name()
+        else:
+            target_or_name = target
+        decomps = self.functions.get(target_or_name, [])
+        return sorted(decomps, key=lambda x: x.is_custom, reverse=True)
+
+    def is_registered(self, target: TorchOp) -> bool:
+        """Returns whether the given op is registered: torch.ops.<namespace>.<op_name>.<overload>.
+
+        Args:
+            target: The PyTorch node callable target.
+
+        Returns:
+            True if the given op is registered, otherwise False.
+        """
+        return bool(self.get_decomps(target))
+
+    def _cleanup_registry_based_on_opset_version(self) -> None:
+        """Pick the implementation with the highest opset version valid until the current opset version."""
+        cleaned_functions = {}
+        for target_or_name, decomps in self.functions.items():
+            # Filter decompositions to only include those with opset_introduced <= opset_version
+            decomps = [d for d in decomps if d.opset_introduced <= self.opset_version]
+
+            # Keep only the decomposition with the highest opset_introduced
+            if decomps:
+                # Find the maximum opset_introduced
+                max_opset = max(d.opset_introduced for d in decomps)
+
+                # Keep all decompositions with the maximum opset_introduced
+                cleaned_functions[target_or_name] = [
+                    d for d in decomps if d.opset_introduced == max_opset
+                ]
+
+        self.functions = cleaned_functions
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(functions={self.functions})"

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_reporting.py

@@ -0,0 +1,207 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import dataclasses
+import re
+from typing import TYPE_CHECKING
+
+from torch.onnx._internal.exporter import _analysis, _registration, _verification
+
+
+if TYPE_CHECKING:
+    import os
+
+    import onnx_ir as ir  # type: ignore[import-untyped]
+
+    import torch
+
+
+@dataclasses.dataclass
+class ExportStatus:
+    # Whether torch.export.export(..., strict=True) succeeds
+    torch_export_strict: bool | None = None
+    # Whether torch.export.export(..., strict=False) succeeds
+    torch_export_non_strict: bool | None = None
+    # Whether torch.export.draft_export() succeeds
+    torch_export_draft_export: bool | None = None
+    # Whether decomposition succeeds
+    decomposition: bool | None = None
+    # Whether ONNX translation succeeds
+    onnx_translation: bool | None = None
+    # Whether ONNX model passes onnx.checker.check_model
+    onnx_checker: bool | None = None
+    # Whether ONNX model runs successfully with ONNX Runtime
+    onnx_runtime: bool | None = None
+    # Whether the output of the ONNX model is accurate
+    output_accuracy: bool | None = None
+
+
+def _status_emoji(status: bool | None) -> str:
+    if status is None:
+        return "⚪"
+    return "✅" if status else "❌"
+
+
+def _format_export_status(status: ExportStatus) -> str:
+    return (
+        f"```\n"
+        f"{_status_emoji(status.torch_export_non_strict)} Obtain model graph with `torch.export.export(..., strict=False)`\n"
+        f"{_status_emoji(status.torch_export_strict)} Obtain model graph with `torch.export.export(..., strict=True)`\n"
+        f"{_status_emoji(status.torch_export_draft_export)} Obtain model graph with `torch.export.draft_export`\n"
+        f"{_status_emoji(status.decomposition)} Decompose operators for ONNX compatibility\n"
+        f"{_status_emoji(status.onnx_translation)} Translate the graph into ONNX\n"
+        f"{_status_emoji(status.onnx_checker)} Run `onnx.checker` on the ONNX model\n"
+        f"{_status_emoji(status.onnx_runtime)} Execute the model with ONNX Runtime\n"
+        f"{_status_emoji(status.output_accuracy)} Validate model output accuracy\n"
+        f"```\n\n"
+    )
+
+
+def _strip_color_from_string(text: str) -> str:
+    # This regular expression matches ANSI escape codes
+    # https://github.com/pytorch/pytorch/blob/9554a9af8788c57e1c5222c39076a5afcf0998ae/torch/_dynamo/utils.py#L2785-L2788
+    ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
+    return ansi_escape.sub("", text)
+
+
+def _format_exported_program(exported_program: torch.export.ExportedProgram) -> str:
+    # Adapted from https://github.com/pytorch/pytorch/pull/128476
+    # to remove colors
+    # Even though we can call graph_module.print_readable directly, since the
+    # colored option was added only recently, we can't guarantee that the
+    # version of PyTorch used by the user has this option. Therefore, we
+    # still call str(ExportedProgram)
+    text = f"```python\n{_strip_color_from_string(str(exported_program))}\n```\n\n"
+    return text
+
+
+def construct_report_file_name(timestamp: str, status: ExportStatus) -> str:
+    # Status could be None. So we need to check for False explicitly.
+    if not (
+        status.torch_export_non_strict
+        or status.torch_export_strict
+        or status.torch_export_draft_export
+    ):
+        # All strategies failed
+        postfix = "pt_export"
+    elif status.decomposition is False:
+        postfix = "decomp"
+    elif status.onnx_translation is False:
+        postfix = "conversion"
+    elif status.onnx_checker is False:
+        postfix = "checker"
+    elif status.onnx_runtime is False:
+        postfix = "runtime"
+    elif status.output_accuracy is False:
+        postfix = "accuracy"
+    elif (
+        status.torch_export_strict is False
+        or status.torch_export_non_strict is False
+        or status.torch_export_draft_export is False
+    ):
+        # Some strategies failed
+        postfix = "strategies"
+    else:
+        postfix = "success"
+    return f"onnx_export_{timestamp}_{postfix}.md"
+
+
+def format_decomp_comparison(
+    pre_decomp_unique_ops: set[str],
+    post_decomp_unique_ops: set[str],
+) -> str:
+    """Format the decomposition comparison result.
+
+    Args:
+        unique_ops_in_a: The unique ops in the first program.
+        unique_ops_in_b: The unique ops in the second program.
+
+    Returns:
+        The formatted comparison result.
+    """
+    return (
+        f"Ops exist only in the ExportedProgram before decomposition: `{sorted(pre_decomp_unique_ops)}`\n\n"
+        f"Ops exist only in the ExportedProgram after decomposition: `{sorted(post_decomp_unique_ops)}`\n"
+    )
+
+
+def format_verification_infos(
+    verification_infos: list[_verification.VerificationInfo],
+) -> str:
+    """Format the verification result.
+
+    Args:
+        verification_infos: The verification result.
+
+    Returns:
+        The formatted verification result.
+    """
+    return "\n".join(
+        f"`{info.name}`: `max_abs_diff={info.max_abs_diff:e}`, `max_rel_diff={info.max_rel_diff:e}`, "
+        f"`abs_diff_hist={info.abs_diff_hist}`, `rel_diff_hist={info.rel_diff_hist}`"
+        for info in verification_infos
+    )
+
+
+def create_torch_export_error_report(
+    filename: str | os.PathLike,
+    formatted_traceback: str,
+    *,
+    export_status: ExportStatus,
+    profile_result: str | None,
+) -> None:
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write("# PyTorch ONNX Conversion Error Report\n\n")
+        f.write(_format_export_status(export_status))
+        f.write("Error message:\n\n")
+        f.write("```pytb\n")
+        f.write(formatted_traceback)
+        f.write("```\n\n")
+        if profile_result is not None:
+            f.write("## Profiling result\n\n")
+            f.write("```\n")
+            f.write(profile_result)
+            f.write("```\n")
+
+
+def create_onnx_export_report(
+    filename: str | os.PathLike,
+    formatted_traceback: str,
+    program: torch.export.ExportedProgram,
+    *,
+    decomp_comparison: str | None = None,
+    export_status: ExportStatus,
+    profile_result: str | None,
+    model: ir.Model | None = None,
+    registry: _registration.ONNXRegistry | None = None,
+    verification_result: str | None = None,
+) -> None:
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write("# PyTorch ONNX Conversion Report\n\n")
+        f.write(_format_export_status(export_status))
+        f.write("## Error messages\n\n")
+        f.write("```pytb\n")
+        f.write(formatted_traceback)
+        f.write("\n```\n\n")
+        f.write("## Exported program\n\n")
+        f.write(_format_exported_program(program))
+        if model is not None:
+            f.write("## ONNX model\n\n")
+            f.write("```python\n")
+            f.write(str(model))
+            f.write("\n```\n\n")
+        f.write("## Analysis\n\n")
+        _analysis.analyze(program, file=f, registry=registry)
+        if decomp_comparison is not None:
+            f.write("\n## Decomposition comparison\n\n")
+            f.write(decomp_comparison)
+            f.write("\n")
+        if verification_result is not None:
+            f.write("\n## Verification results\n\n")
+            f.write(verification_result)
+            f.write("\n")
+        if profile_result is not None:
+            f.write("\n## Profiling result\n\n")
+            f.write("```\n")
+            f.write(profile_result)
+            f.write("```\n")

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_schemas.py

@@ -0,0 +1,332 @@
+# mypy: allow-untyped-defs
+"""Helpers for constructing ONNX operator signatures from Python functions."""
+
+from __future__ import annotations
+
+import collections.abc
+import inspect
+import logging
+import types
+import typing
+from collections.abc import Sequence
+from typing import Any, Optional, TypeVar, Union
+
+from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript
+
+
+logger = logging.getLogger(__name__)
+
+
+# Map from python type to corresponding ONNX AttributeProto type
+_PY_TYPE_TO_ATTR_TYPE = {
+    float: ir.AttributeType.FLOAT,
+    int: ir.AttributeType.INT,
+    str: ir.AttributeType.STRING,
+    bool: ir.AttributeType.INT,
+    ir.Tensor: ir.AttributeType.TENSOR,
+    ir.TensorProtocol: ir.AttributeType.TENSOR,
+    ir.Graph: ir.AttributeType.GRAPH,
+    ir.GraphProtocol: ir.AttributeType.GRAPH,
+}
+
+# Map from python type to corresponding ONNX AttributeProto type,
+# for repeated (i.e., list of) values
+_LIST_TYPE_TO_ATTR_TYPE = {
+    float: ir.AttributeType.FLOATS,
+    int: ir.AttributeType.INTS,
+    str: ir.AttributeType.STRINGS,
+    bool: ir.AttributeType.INTS,
+    ir.Tensor: ir.AttributeType.TENSORS,
+    ir.TensorProtocol: ir.AttributeType.TENSORS,
+    ir.Graph: ir.AttributeType.GRAPHS,
+    ir.GraphProtocol: ir.AttributeType.GRAPHS,
+}
+
+_ALL_VALUE_TYPES = (
+    {ir.TensorType(dtype) for dtype in ir.DataType}
+    | {ir.SequenceType(ir.TensorType(dtype)) for dtype in ir.DataType}
+    | {ir.OptionalType(ir.TensorType(dtype)) for dtype in ir.DataType}
+)
+
+# TypeAnnotationValue represents the (value of) valid type-annotations recognized
+# by ONNX Script. Currently, it supports
+# - float, int, str (primitive attribute types)
+# - Sequence[float], Sequence[int], Sequence[str] (attribute types)
+# - Tensor types
+# - Sequence[Tensor] types
+# - Union of above 2
+# - TypeVars with above bounds
+# - Above types with annotation attached
+TypeAnnotationValue = Any
+
+
+def _is_optional(type_: type) -> bool:
+    """Returns whether a type_ is an Optional."""
+    origin_type = typing.get_origin(type_)
+    if origin_type is Union and type(None) in typing.get_args(type_):
+        # Python < 3.10
+        return True
+    if origin_type is Optional:
+        # Python >= 3.10
+        return True
+    if (
+        hasattr(types, "UnionType")
+        and origin_type is types.UnionType
+        and type(None) in typing.get_args(type_)
+    ):
+        # Python >= 3.10
+        return True
+    return False
+
+
+def _get_attr_type(type_: type) -> ir.AttributeType:
+    """Obtain the type of the attribute from a Python class."""
+    try:
+        if type_ in _PY_TYPE_TO_ATTR_TYPE:
+            return _PY_TYPE_TO_ATTR_TYPE[type_]
+        origin_type = typing.get_origin(type_)
+        if origin_type is None:
+            return ir.AttributeType.UNDEFINED
+        if origin_type in (
+            collections.abc.Sequence,
+            Sequence,
+            list,
+            list,
+            tuple,
+            tuple,
+        ):
+            inner_type = typing.get_args(type_)[0]
+            if inner_type in _LIST_TYPE_TO_ATTR_TYPE:
+                return _LIST_TYPE_TO_ATTR_TYPE[inner_type]
+    except TypeError:
+        logger.warning("TypeError when checking %s.", type_, exc_info=True)
+    return ir.AttributeType.UNDEFINED
+
+
+def _get_type_constraint_name(type_: TypeAnnotationValue) -> str | None:
+    """Returns the name of the type constraint for a given type annotation.
+
+    Args:
+        type_: A Python type.
+
+    Returns:
+        The name of the type constraint if it is a TypeVar.
+        - Prefixes the name with "Sequence_" if the type annotation is a Sequence[].
+    """
+    if isinstance(type_, TypeVar):
+        return type_.__name__
+    if _is_optional(type_):
+        subtypes = typing.get_args(type_)
+        for subtype in subtypes:
+            if subtype is type(None):
+                continue
+            type_param_name = _get_type_constraint_name(subtype)
+            return type_param_name if type_param_name else None
+    origin_type = typing.get_origin(type_)
+    if isinstance(origin_type, type) and issubclass(origin_type, Sequence):
+        subtypes = typing.get_args(type_)
+        type_param_name = _get_type_constraint_name(subtypes[0])
+        return f"Sequence_{type_param_name}" if type_param_name else None
+    return None
+
+
+def _get_allowed_types_from_type_annotation(
+    type_: TypeAnnotationValue,
+) -> set[ir.TypeProtocol]:
+    """Obtain the allowed types from a type annotation."""
+    if type_ is onnxscript.onnx_types.TensorType:
+        # Any tensor type
+        return {ir.TensorType(dtype) for dtype in ir.DataType}
+
+    allowed_types: set[ir.TypeProtocol]
+
+    if isinstance(type_, TypeVar):
+        allowed_types = set()
+        if constraints := type_.__constraints__:
+            for constraint in constraints:
+                allowed_types.update(
+                    _get_allowed_types_from_type_annotation(constraint)
+                )
+        else:
+            bound = type_.__bound__
+            if bound is None:
+                allowed_types = _ALL_VALUE_TYPES  # type: ignore[assignment]
+            else:
+                allowed_types.update(_get_allowed_types_from_type_annotation(bound))
+        return allowed_types
+    if hasattr(type_, "dtype"):
+        # A single tensor type like INT64, FLOAT, etc.
+        return {ir.TensorType(ir.DataType(type_.dtype))}
+    if _is_optional(type_):
+        allowed_types = set()
+        subtypes = typing.get_args(type_)
+        for subtype in subtypes:
+            if subtype is type(None):
+                continue
+            allowed_types.update(_get_allowed_types_from_type_annotation(subtype))
+        # NOTE: We do not consider dynamic optional types like optional(float) because they are not very useful.
+        return allowed_types
+
+    origin_type = typing.get_origin(type_)
+    if origin_type is Union:
+        allowed_types = set()
+        subtypes = typing.get_args(type_)
+        for subtype in subtypes:
+            if subtype is type(None):
+                raise AssertionError(
+                    "Union should not contain None type because it is handled by _is_optional."
+                )
+            allowed_types.update(_get_allowed_types_from_type_annotation(subtype))
+        return allowed_types
+
+    if isinstance(origin_type, type) and issubclass(origin_type, Sequence):
+        subtypes = typing.get_args(type_)
+        return {
+            ir.SequenceType(t)
+            for t in _get_allowed_types_from_type_annotation(subtypes[0])
+        }
+
+    # Allow everything by default
+    return _ALL_VALUE_TYPES  # type: ignore[return-value]
+
+
+def op_signature_from_function(
+    func,
+    domain: str,
+    name: str | None = None,
+    overload: str = "",
+    *,
+    since_version: int = 1,
+) -> ir.schemas.OpSignature:
+    """Produce an OpSignature from a function using type annotation."""
+
+    py_signature = inspect.signature(func)
+    # Not using inspect.get_annotations because typing.get_type_hints seems to handle more cases
+    # https://github.com/python/cpython/issues/102405
+    type_hints = typing.get_type_hints(func)
+
+    params: list[ir.schemas.Parameter | ir.schemas.AttributeParameter] = []
+    # Create a mapping from type to a unique name
+    type_constraints: dict[str, ir.schemas.TypeConstraintParam] = {}
+
+    for param in py_signature.parameters.values():
+        if param.name not in type_hints:
+            logger.debug(
+                "Missing annotation for parameter '%s' from %s. Treating as an Input.",
+                param.name,
+                py_signature,
+            )
+            type_constraint = ir.schemas.TypeConstraintParam.any_value(
+                f"T_{param.name}"
+            )
+            type_constraints[param.name] = type_constraint
+            kwargs: dict[str, Any] = {}
+            if param.default is not inspect.Parameter.empty:
+                kwargs["default"] = param.default
+            params.append(
+                ir.schemas.Parameter(
+                    name=param.name,
+                    type_constraint=type_constraint,
+                    required=param.default is inspect.Parameter.empty,
+                    # TODO: Handle variadic
+                    variadic=False,
+                    **kwargs,
+                )
+            )
+        else:
+            type_ = type_hints[param.name]
+            if (attr_type := _get_attr_type(type_)) != ir.AttributeType.UNDEFINED:
+                # Construct the default attribute
+                if param.default is not inspect.Parameter.empty:
+                    # TODO: Use ir_convenience instead to handle int as float
+                    default = ir.Attr(param.name, attr_type, param.default)
+                else:
+                    default = None
+                params.append(
+                    ir.schemas.AttributeParameter(
+                        name=param.name,
+                        type=attr_type,
+                        required=param.default is inspect.Parameter.empty,
+                        default=default,
+                    )
+                )
+            else:
+                # Obtain the type constraint from the type annotation
+
+                # 1. Get a type constraint name from the type annotation
+                # If the type annotation is a TypeVar or Optional[TypeVar], get its name
+                # Otherwise, name it T_{param.name}
+                type_constraint_name = _get_type_constraint_name(type_)
+                if type_constraint_name is None:
+                    type_constraint_name = f"T_{param.name}"
+
+                # 2. If the type constraint param is already initialized, use it
+                if type_constraint_name in type_constraints:
+                    type_constraint = type_constraints[type_constraint_name]
+                else:
+                    # 3. Otherwise, create a new TypeConstraintParam
+                    type_constraint = ir.schemas.TypeConstraintParam(
+                        name=type_constraint_name,
+                        allowed_types=_get_allowed_types_from_type_annotation(type_),
+                    )
+                    type_constraints[type_constraint_name] = type_constraint
+                # 4. Create Parameter
+                kwargs: dict[str, Any] = {}
+                if param.default is not inspect.Parameter.empty:
+                    kwargs["default"] = param.default
+                params.append(
+                    ir.schemas.Parameter(
+                        name=param.name,
+                        type_constraint=type_constraint,
+                        required=param.default is inspect.Parameter.empty,
+                        # TODO: Handle variadic
+                        variadic=False,
+                        **kwargs,
+                    )
+                )
+
+    return_type = type_hints.get("return")
+
+    outputs = []
+    if return_type is None:
+        # No returns
+        pass
+    else:
+        if typing.get_origin(return_type) is tuple:
+            # Multiple returns
+            return_types = typing.get_args(return_type)
+        else:
+            return_types = [return_type]  # type: ignore[assignment]
+
+        for i, return_type_i in enumerate(return_types):
+            if (
+                return_param_name := _get_type_constraint_name(return_type_i)
+            ) in type_constraints:
+                # pyrefly: ignore [bad-index]
+                type_constraint = type_constraints[return_param_name]
+            else:
+                return_param_name = f"TReturn{i}"
+                type_constraint = ir.schemas.TypeConstraintParam(
+                    name=return_param_name,
+                    allowed_types=_get_allowed_types_from_type_annotation(
+                        return_type_i
+                    ),
+                )
+                type_constraints[return_param_name] = type_constraint
+            outputs.append(
+                ir.schemas.Parameter(
+                    name=return_param_name,
+                    type_constraint=type_constraint,
+                    required=True,
+                    variadic=False,
+                )
+            )
+
+    return ir.schemas.OpSignature(
+        domain=domain,
+        name=name or func.__name__,
+        overload=overload,
+        params=params,
+        outputs=outputs,
+        since_version=since_version,
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_tensors.py

@@ -0,0 +1,100 @@
+"""Subclass of ir.Value that supports Python operators."""
+
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+
+
+if TYPE_CHECKING:
+    import onnxscript
+
+
+class SymbolicTensor(ir.Value):
+    """A subclass of ir.Value that supports Python operators."""
+
+    def __init__(
+        self,
+        opset: onnxscript.values.Opset,
+        name: str | None = None,
+        shape: ir.Shape | None = None,
+        type: ir.TypeProtocol | None = None,
+        doc_string: str | None = None,
+        const_value: ir.TensorProtocol | None = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            shape=shape,
+            type=type,
+            doc_string=doc_string,
+            const_value=const_value,
+        )
+        self._opset = opset
+
+    @property
+    def rank(self) -> int | None:
+        if self.shape is None:
+            return None
+        return len(self.shape)
+
+    # TODO: Implement indexing
+
+    def __mod__(self, other):
+        if self.dtype in {
+            ir.DataType.FLOAT,
+            ir.DataType.DOUBLE,
+            ir.DataType.FLOAT16,
+            ir.DataType.BFLOAT16,
+        }:
+            return self._opset.Mod(self, other, fmod=1)
+        return self._opset.Mod(self, other)
+
+    def __ne__(self, other):
+        return self._opset.Not(self._opset.Equal(self, other))
+
+    def __neg__(self):
+        return self._opset.Neg(self)
+
+    def __add__(self, other):
+        return self._opset.Add(self, other)
+
+    def __radd__(self, other):
+        return self._opset.Add(other, self)
+
+    def __rand__(self, other):
+        return self._opset.And(other, self)
+
+    def __mul__(self, other):
+        return self._opset.Mul(self, other)
+
+    def __rmul__(self, other):
+        return self._opset.Mul(other, self)
+
+    def __matmul__(self, other):
+        return self._opset.MatMul(self, other)
+
+    def __pow__(self, other):
+        return self._opset.Pow(self, other)
+
+    def __sub__(self, other):
+        return self._opset.Sub(self, other)
+
+    def __rsub__(self, other):
+        return self._opset.Sub(other, self)
+
+    def __truediv__(self, other):
+        return self._opset.Div(self, other)
+
+    def __lt__(self, other):
+        return self._opset.Less(self, other)
+
+    def __le__(self, other):
+        return self._opset.LessOrEqual(self, other)
+
+    def __ge__(self, other):
+        return self._opset.GreaterOrEqual(self, other)
+
+    def __gt__(self, other):
+        return self._opset.Greater(self, other)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_testing.py

@@ -0,0 +1,102 @@
+"""Test utilities for ONNX export."""
+
+from __future__ import annotations
+
+
+__all__ = ["assert_onnx_program"]
+
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+from torch.utils import _pytree
+
+
+if TYPE_CHECKING:
+    from torch.onnx._internal.exporter import _onnx_program
+
+
+def assert_onnx_program(
+    program: _onnx_program.ONNXProgram,
+    *,
+    rtol: float | None = None,
+    atol: float | None = None,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    strategy: str | None = "TorchExportNonStrictStrategy",
+    backend: Literal["onnxruntime", "reference"] = "onnxruntime",
+) -> None:
+    """Assert that the ONNX model produces the same output as the PyTorch ExportedProgram.
+
+    Args:
+        program: The ``ONNXProgram`` to verify.
+        rtol: Relative tolerance.
+        atol: Absolute tolerance.
+        args: The positional arguments to pass to the program.
+            If None, the default example inputs in the ExportedProgram will be used.
+        kwargs: The keyword arguments to pass to the program.
+            If None, the default example inputs in the ExportedProgram will be used.
+        strategy: Assert the capture strategy used to export the program. Values can be
+            class names like "TorchExportNonStrictStrategy".
+            If None, the strategy is not asserted.
+        backend: The backend to use for evaluating the ONNX program.
+            Supported values are "onnxruntime" and "reference".
+    """
+    if strategy is not None:
+        if program._capture_strategy != strategy:
+            raise ValueError(
+                f"Expected strategy '{strategy}' is used to capture the exported program, "
+                f"but got '{program._capture_strategy}'."
+            )
+    exported_program = program.exported_program
+    if exported_program is None:
+        raise ValueError(
+            "The ONNXProgram does not contain an ExportedProgram. "
+            "To verify the ONNX program, initialize ONNXProgram with an ExportedProgram, "
+            "or assign the ExportedProgram to the ONNXProgram.exported_program attribute."
+        )
+    if args is None and kwargs is None:
+        # User did not provide example inputs, use the default example inputs
+        if exported_program.example_inputs is None:
+            raise ValueError(
+                "No example inputs provided and the exported_program does not contain example inputs. "
+                "Please provide arguments to verify the ONNX program."
+            )
+        args, kwargs = exported_program.example_inputs
+    if args is None:
+        args = ()
+    if kwargs is None:
+        kwargs = {}
+    torch_module = exported_program.module()
+    torch_outputs, _ = _pytree.tree_flatten(torch_module(*args, **kwargs))
+    # ONNX outputs are always real, so we need to convert torch complex outputs to real representations
+    torch_outputs_adapted = []
+    for output in torch_outputs:
+        # ONNX graph does not support None outputs, so we skip them
+        if output is None:
+            continue
+        if not isinstance(output, torch.Tensor):
+            torch_outputs_adapted.append(torch.tensor(output))
+        elif torch.is_complex(output):
+            torch_outputs_adapted.append(torch.view_as_real(output))
+        else:
+            torch_outputs_adapted.append(output)
+
+    # Obtain the ONNX outputs using the specified backend
+    if backend == "onnxruntime":
+        onnx_outputs = program(*args, **kwargs)
+    elif backend == "reference":
+        onnx_outputs = program.call_reference(*args, **kwargs)
+    else:
+        raise ValueError(
+            f"Unsupported backend '{backend}'. Supported backends are 'onnxruntime' and 'reference'."
+        )
+
+    # TODO(justinchuby): Include output names in the error message
+    torch.testing.assert_close(
+        tuple(onnx_outputs),
+        tuple(torch_outputs_adapted),
+        rtol=rtol,
+        atol=atol,
+        equal_nan=True,
+        check_device=False,
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py

@@ -0,0 +1,75 @@
+"""Typings for function definitions."""
+
+from __future__ import annotations
+
+from typing import TypeVar, Union
+
+from onnxscript import (
+    BFLOAT16,
+    BOOL,
+    COMPLEX128,
+    COMPLEX64,
+    DOUBLE,
+    FLOAT,
+    FLOAT16,
+    INT16,
+    INT32,
+    INT64,
+    INT8,
+    STRING,
+    UINT8,
+)
+
+
+# NOTE: We do not care about unsigned types beyond UINT8 because PyTorch does not us them.
+# More detail can be found: https://pytorch.org/docs/stable/tensors.html
+
+TensorType = Union[  # noqa: UP007
+    BFLOAT16,
+    BOOL,
+    COMPLEX64,
+    COMPLEX128,
+    DOUBLE,
+    FLOAT,
+    FLOAT16,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    UINT8,
+]
+_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]  # noqa: UP007
+IntType = Union[INT8, INT16, INT32, INT64]  # noqa: UP007
+RealType = Union[  # noqa: UP007
+    BFLOAT16,
+    FLOAT16,
+    FLOAT,
+    DOUBLE,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+]
+
+TTensor = TypeVar("TTensor", bound=TensorType)
+# Duplicate TTensor for inputs/outputs that accept the same set of types as TTensor
+# but do not constrain the type to be the same as the other inputs/outputs
+TTensor2 = TypeVar("TTensor2", bound=TensorType)
+TTensorOrString = TypeVar("TTensorOrString", bound=Union[TensorType, STRING])  # noqa: UP007
+TFloat = TypeVar("TFloat", bound=_FloatType)
+TFloatOrUInt8 = TypeVar(
+    "TFloatOrUInt8",
+    bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8],  # noqa: UP007
+)
+TInt = TypeVar("TInt", bound=IntType)
+TReal = TypeVar("TReal", bound=RealType)
+TRealUnlessInt16OrInt8 = TypeVar(
+    "TRealUnlessInt16OrInt8",
+    bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64],  # noqa: UP007
+)
+TRealUnlessFloat16OrInt8 = TypeVar(
+    "TRealUnlessFloat16OrInt8",
+    bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64],  # noqa: UP007
+)
+TRealOrUInt8 = TypeVar("TRealOrUInt8", bound=Union[RealType, UINT8])  # noqa: UP007
+TFloatHighPrecision = TypeVar("TFloatHighPrecision", bound=Union[FLOAT, DOUBLE])  # noqa: UP007

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/_torchlib_registry.py

@@ -0,0 +1,95 @@
+"""Registry for aten functions."""
+
+from __future__ import annotations
+
+
+__all__ = ["onnx_impl", "get_torchlib_ops"]
+
+import logging
+from collections.abc import Callable, Sequence
+from typing import Any, TypeVar
+from typing_extensions import ParamSpec
+
+import onnxscript
+
+import torch
+from torch.onnx._internal.exporter import _constants, _registration
+
+
+# Use ParamSpec for better type preservation instead of bound Callable TypeVar
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+logger = logging.getLogger("__name__")
+
+
+_registry: list[_registration.OnnxDecompMeta] = []
+
+
+def onnx_impl(
+    target: _registration.TorchOp | tuple[_registration.TorchOp, ...],
+    *,
+    trace_only: bool = False,
+    complex: bool = False,
+    opset_introduced: int = 18,
+    no_compile: bool = False,
+    private: bool = False,
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+    """Register an ONNX implementation of a torch op."""
+
+    if isinstance(target, torch._ops.OpOverloadPacket):
+        raise TypeError(
+            f"Target '{target}' should be provided as an OpOverload instead of an "
+            "OpOverloadPacket. You can get the default overload with "
+            "<op>.default"
+        )
+
+    def wrapper(
+        func: Callable[_P, _R],
+    ) -> Callable[_P, _R]:
+        processed_func: Any
+        if no_compile:
+            processed_func = func
+        else:
+            torchlib_opset = onnxscript.values.Opset(
+                domain=_constants.TORCHLIB_DOMAIN, version=1
+            )
+
+            if not trace_only:
+                # Compile the function
+                processed_func = onnxscript.script(opset=torchlib_opset)(func)
+            else:
+                processed_func = onnxscript.TracedOnnxFunction(torchlib_opset, func)
+
+        if not private:
+            # TODO(justinchuby): Simplify the logic and remove the private attribute
+            # Skip registration if private
+            if not isinstance(target, Sequence):
+                targets = (target,)
+            else:
+                targets = target  # type: ignore[assignment]
+
+            for t in targets:
+                _registry.append(
+                    _registration.OnnxDecompMeta(
+                        onnx_function=processed_func,
+                        fx_target=t,
+                        signature=None,
+                        is_complex=complex,
+                        opset_introduced=opset_introduced,
+                        skip_signature_inference=no_compile,
+                    )
+                )
+        return processed_func  # type: ignore[return-value]
+
+    return wrapper
+
+
+def get_torchlib_ops() -> tuple[_registration.OnnxDecompMeta, ...]:
+    # Trigger op registration
+    from torch.onnx._internal.exporter._torchlib import ops
+
+    del ops
+    if len(_registry) == 0:
+        raise AssertionError("_registry must not be empty")
+    return tuple(_registry)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+
+__all__ = ["core", "hop", "nn", "symbolic", "symops"]
+
+from torch.onnx._internal.exporter._torchlib.ops import core, hop, nn, symbolic, symops

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/core.py

@@ -0,0 +1,47 @@
+"""torch.ops.aten operators under the `core` module."""
+# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index"
+# pyrefly: ignore-errors
+# ruff: noqa: TCH001,TCH002
+
+from __future__ import annotations
+
+import operator
+
+from onnxscript.onnx_opset import opset18 as op
+
+import torch
+from torch.onnx._internal.exporter._torchlib._tensor_typing import TReal, TRealOrUInt8
+from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl
+
+
+aten = torch.ops.aten
+
+
+@onnx_impl((aten.abs.default, operator.abs), trace_only=True)
+def aten_abs(self: TRealOrUInt8) -> TRealOrUInt8:
+    """abs(Tensor self) -> Tensor"""
+
+    return op.Abs(self)
+
+
+@onnx_impl(aten.abs.default, complex=True, trace_only=True)
+def aten_abs_complex(self: TRealOrUInt8) -> TRealOrUInt8:
+    """abs(Tensor self) -> Tensor"""
+
+    return op.ReduceL2(self, [-1], keepdims=False)
+
+
+@onnx_impl((aten.add.Tensor, aten.add.Scalar, operator.add), trace_only=True)
+def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
+    """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"""
+    if alpha != 1.0:
+        alpha = op.CastLike(alpha, other)
+        other = op.Mul(other, alpha)
+    return op.Add(self, other)
+
+
+@onnx_impl((aten.add.Tensor, aten.add.Scalar), trace_only=True, complex=True)
+def aten_add_complex(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
+    """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"""
+
+    return aten_add(self, other, alpha=alpha)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/hop.py

@@ -0,0 +1,370 @@
+"""Implementation for higher-order operators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+from torch.onnx._internal.exporter import _core
+from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+def call_op(
+    op_type: str,
+    *args: ir.Value | None,
+    _num_outputs: int = 1,
+    _domain: str = "",
+    **kwargs: int | float | str | bool | ir.Graph | ir.TensorProtocol | Sequence[int],
+) -> Sequence[ir.Value]:
+    """Call an operator with the given arguments and keyword arguments.
+
+    Arguments are always inputs, while keyword arguments are attributes.
+    """
+    # This is a wrapper around the IR node creation that hooks into the _builder.OpRecorder
+    # tracer so that all nodes created are recorded the same way as if we were to use
+    # onnxscript ops directly.
+    from onnxscript.ir import convenience as ir_convenience
+
+    if _core.current_tracer is None:
+        raise AssertionError("current_tracer must be non-None")
+    tracer = _core.current_tracer
+
+    inputs = list(args)
+
+    # If final inputs are None, strip them from the node inputs
+    for input in reversed(inputs):
+        if input is not None:
+            break
+        inputs.pop()
+
+    # Construct and filter out None attributes
+    attributes = [
+        attr
+        for attr in ir_convenience.convert_attributes(kwargs)
+        if attr.value is not None  # type: ignore[union-attr]
+    ]
+    tracer.nodes.append(
+        node := ir.Node(
+            _domain,
+            op_type,
+            inputs=inputs,
+            attributes=attributes,
+            num_outputs=_num_outputs,
+            version=tracer.opset.version,
+        )
+    )
+    return node.outputs
+
+
+@onnx_impl(torch.ops.higher_order.cond, no_compile=True)
+def higher_order_cond(
+    cond: ir.Value,
+    true_func: ir.Function,
+    false_func: ir.Function,
+    inputs: Sequence[ir.Value],
+) -> Sequence[ir.Value]:
+    then_node = ir.Node(
+        true_func.domain, true_func.name, inputs, num_outputs=len(true_func.outputs)
+    )
+    else_node = ir.Node(
+        false_func.domain, false_func.name, inputs, num_outputs=len(false_func.outputs)
+    )
+
+    # ONNX Runtime complains about duplicate output names if we don't rename them.
+    # But the doesn't seem to be an actual violation of SSA form without renaming.
+    for func_out, out in zip(true_func.outputs, then_node.outputs):
+        out.name = f"{func_out.name}_{true_func.name}"
+    for func_out, out in zip(false_func.outputs, else_node.outputs):
+        out.name = f"{func_out.name}_{false_func.name}"
+
+    return call_op(
+        "If",
+        cond,
+        _num_outputs=len(true_func.outputs),
+        then_branch=ir.Graph(
+            (), then_node.outputs, nodes=[then_node], name=true_func.name
+        ),
+        else_branch=ir.Graph(
+            (), else_node.outputs, nodes=[else_node], name=false_func.name
+        ),
+    )
+
+
+@onnx_impl(torch.ops.higher_order.scan, no_compile=True)
+def higher_order_scan(
+    body_func: ir.Function,
+    scan_inits: Sequence[ir.Value],
+    scan_inputs: Sequence[ir.Value],
+    additional_inputs: Sequence[ir.Value] | None,
+    reverse: bool = False,
+) -> Sequence[ir.Value]:
+    """https://github.com/pytorch/pytorch/blob/66ac724b56e6c37a534f3e066423ef2f41d7477f/torch/_higher_order_ops/scan.py#L109"""
+    subgraph_inputs = [
+        *[
+            ir.Value(
+                name=f"{inp.name}_{body_func.name}__subgraph_in",
+                shape=inp.shape,
+                type=ir.TensorType(inp.dtype),  # type: ignore[arg-type]
+            )
+            for inp in scan_inits
+        ],
+        *[
+            ir.Value(
+                name=f"{inp.name}_{body_func.name}__subgraph_in",
+                # The iterated element passed to the body subgraph does not have a sequence axis.
+                # It will have a rank one less than the rank of the corresponding scan_input.
+                shape=ir.Shape(inp.shape[1:]),  # type: ignore[index]
+                type=ir.TensorType(inp.dtype),  # type: ignore[arg-type]
+            )
+            for inp in scan_inputs
+        ],
+    ]
+    # The one and only node in the Scan subgraph that calls the body_func
+    body_node = ir.Node(
+        body_func.domain,
+        body_func.name,
+        [
+            *subgraph_inputs,
+            *(additional_inputs or []),
+        ],
+        num_outputs=len(body_func.outputs),
+    )
+
+    # ONNX Runtime complains about duplicate output names if we don't rename them.
+    # But the doesn't seem to be an actual violation of SSA form without renaming.
+    for func_out, out in zip(body_func.outputs, body_node.outputs):
+        out.name = f"{func_out.name}_{body_func.name}"
+
+    n_outputs = len(body_func.outputs) - len(scan_inits)
+    return call_op(
+        "Scan",
+        *scan_inits,
+        *scan_inputs,
+        _num_outputs=len(body_func.outputs),
+        body=ir.Graph(
+            subgraph_inputs,
+            body_node.outputs,
+            nodes=[body_node],
+            name=body_func.name,
+        ),
+        num_scan_inputs=len(scan_inputs),
+        scan_input_directions=[(1 if reverse else 0) for _ in scan_inputs],
+        scan_output_directions=[(1 if reverse else 0) for _ in range(n_outputs)],
+    )
+
+
+@onnx_impl(torch.ops.higher_order.while_loop, no_compile=True)
+def higher_order_while_loop(
+    cond_func: ir.Function,
+    body_func: ir.Function,
+    carried_inputs: Sequence[ir.Value | int | float],
+    additional_inputs: Sequence[ir.Value],
+) -> Sequence[ir.Value]:
+    """Implementation of while_loop using ONNX Loop operator.
+
+    The ONNX Loop operator implements a generic looping construct with the signature:
+    Loop(M, cond, v_initial) -> (v_final_and_scan_outputs)
+
+    For while_loop, we use:
+    - M: None (no trip count limit)
+    - cond: initial condition value
+    - v_initial: carried_inputs (loop-carried dependencies)
+
+    The body subgraph takes:
+    - iteration_num (int): current iteration number
+    - condition_in (bool): loop continuation condition from previous iteration
+    - loop_carried_dependencies: the carried values
+    - additional_inputs: any additional inputs (constants/parameters)
+
+    The body subgraph returns:
+    - condition_out (bool): whether to continue looping
+    - loop_carried_dependencies: updated carried values
+    """
+
+    # Create subgraph inputs for the Loop body
+    # ONNX Loop body signature: (iter_num, cond_in, loop_carried_deps..., additional_inputs...)
+
+    # Start subgraph construction
+    subgraph_carried_inputs = []
+
+    for i, inp in enumerate(carried_inputs):
+        if isinstance(inp, ir.Value):
+            subgraph_carried_inputs.append(
+                ir.Value(
+                    name=f"{inp.name}_{body_func.name}__subgraph_in",
+                    shape=inp.shape,
+                    type=ir.TensorType(inp.dtype),  # type: ignore[arg-type]
+                )
+            )
+        elif isinstance(inp, int):
+            subgraph_carried_inputs.append(
+                ir.Value(
+                    name=f"carried_input_{i}_{body_func.name}__subgraph_in",
+                    shape=ir.Shape([]),
+                    type=ir.TensorType(ir.DataType.INT64),
+                )
+            )
+        elif isinstance(inp, float):
+            subgraph_carried_inputs.append(
+                ir.Value(
+                    name=f"carried_input_{i}_{body_func.name}__subgraph_in",
+                    shape=ir.Shape([]),
+                    type=ir.TensorType(ir.DataType.FLOAT),
+                )
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported type for carried input: {type(inp)} ({inp}). "
+                "Expected ir.Value, int, or float."
+            )
+
+    subgraph_inputs = [
+        # Iteration number (int scalar, unused)
+        ir.Value(
+            name=f"iter_num_{body_func.name}",
+            shape=ir.Shape([]),
+            type=ir.TensorType(ir.DataType.INT64),
+        ),
+        # Condition input (bool scalar, unused)
+        ir.Value(
+            name=f"cond_in_{body_func.name}",
+            shape=ir.Shape([]),
+            type=ir.TensorType(ir.DataType.BOOL),
+        ),
+        # Loop-carried dependencies
+        *subgraph_carried_inputs,
+    ]
+
+    # Create the combined body function that handles both condition and body logic
+    # First, call the body function with the same inputs
+    body_node = ir.Node(
+        body_func.domain,
+        body_func.name,
+        [
+            *subgraph_carried_inputs,  # carried inputs
+            *additional_inputs,
+        ],
+        num_outputs=len(body_func.outputs),  # carried inputs
+    )
+
+    # Then call the condition function with carried inputs + additional inputs
+    cond_node = ir.Node(
+        cond_func.domain,
+        cond_func.name,
+        [
+            *body_node.outputs,  # updated carried inputs from body
+            *additional_inputs,
+        ],
+        num_outputs=len(cond_func.outputs),
+    )
+
+    if len(cond_func.outputs) != 1:
+        raise AssertionError("Condition function must return a single boolean value.")
+
+    # ONNX Runtime complains about duplicate output names if we don't rename them
+    for func_out, out in zip(body_func.outputs, body_node.outputs):
+        out.name = f"{func_out.name}_{body_func.name}"
+    for func_out, out in zip(cond_func.outputs, cond_node.outputs):
+        out.name = f"{func_out.name}_{cond_func.name}"
+
+    # The Loop body must return: (cond_out, loop_carried_deps...)
+    # We use the condition output and the body outputs
+    loop_body_outputs = [
+        cond_node.outputs[0],  # condition output (bool)
+        *body_node.outputs,  # updated carried inputs
+    ]
+
+    body_graph = ir.Graph(
+        subgraph_inputs,
+        loop_body_outputs,
+        nodes=[body_node, cond_node],
+        name=f"{body_func.name}_loop_body",
+    )
+
+    # End subgraph construction
+
+    carried_inputs_values: list[ir.Value] = []
+    for inp in carried_inputs:
+        if isinstance(inp, ir.Value):
+            carried_inputs_values.append(inp)
+        elif isinstance(inp, int):
+            const = call_op("Constant", value=ir.tensor(inp))[0]
+            carried_inputs_values.append(const)
+        elif isinstance(inp, float):
+            const = call_op("Constant", value=ir.tensor(inp))[0]
+            carried_inputs_values.append(const)
+        else:
+            raise NotImplementedError(
+                f"Unsupported type for carried input: {type(inp)} ({inp}). "
+                "Expected ir.Value, int, or float."
+            )
+
+    # Get initial condition by calling cond_func with initial inputs
+    initial_outputs = call_op(
+        cond_func.name,
+        *carried_inputs_values,
+        *additional_inputs,
+        _num_outputs=len(cond_func.outputs),
+        _domain=cond_func.domain,
+    )
+
+    if len(initial_outputs) != 1:
+        raise AssertionError("Condition function must return a single boolean value.")
+
+    # Create the Loop operator call
+    # Loop(M, cond, v_initial) where M is empty (no trip count limit)
+    loop_outputs = call_op(
+        "Loop",
+        # M (trip count) - empty string means no limit
+        None,
+        # cond - initial condition
+        initial_outputs[0],
+        # v_initial - carried inputs (loop-carried dependencies)
+        *carried_inputs_values,
+        _num_outputs=len(carried_inputs_values),
+        body=body_graph,
+    )
+
+    return loop_outputs
+
+
+@onnx_impl(torch.ops.higher_order.invoke_subgraph, no_compile=True)
+def higher_order_invoke_subgraph(
+    subgraph: ir.Function,
+    identifier: str | None,
+    *operands: ir.Value,
+) -> Sequence[ir.Value]:
+    """Export invoke_subgraph HOP by creating a direct function call.
+
+    This preserves the function as a separate entity in the ONNX graph
+    instead of inlining it, which is the purpose of invoke_subgraph.
+
+    Note: The onnxscript optimizer should be configured to not inline functions
+    created by invoke_subgraph to preserve the intended structure.
+
+    Args:
+        subgraph: The function to invoke
+        identifier: Optional identifier for the subgraph (used for caching in PyTorch,
+            not needed for ONNX export as the function reference provides all necessary information)
+        *operands: Input values to pass to the function
+
+    Returns:
+        Sequence of output values from the function call
+    """
+    # This key can be used by downstream to avoid inlining
+    subgraph.metadata_props["pkg.torch.ops.higher_order.invoke_subgraph.identifier"] = (
+        str(identifier)
+    )
+
+    # Create the function call node
+    return call_op(
+        subgraph.name,
+        *operands,
+        _num_outputs=len(subgraph.outputs),
+        _domain=subgraph.domain,
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/nn.py

@@ -0,0 +1,377 @@
+"""torch.ops.aten operators under the `core` module."""
+# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index"
+# pyrefly: ignore-errors
+# ruff: noqa: TC001,TC002
+# flake8: noqa: B950
+
+from __future__ import annotations
+
+from typing import Sequence, TYPE_CHECKING  # noqa: UP035
+
+from onnxscript.onnx_opset import (  # type: ignore[attr-defined]
+    opset20 as op20,
+    opset21 as op21,
+    opset23 as op23,
+)
+
+import torch
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+from torch.onnx._internal.exporter._torchlib._tensor_typing import TFloat, TReal
+from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl
+
+
+if TYPE_CHECKING:
+    from onnxscript.values import Opset
+
+aten = torch.ops.aten
+
+
+@onnx_impl(aten.gelu.default, trace_only=True, opset_introduced=20)
+def aten_gelu_opset20(
+    self: TReal,
+    approximate: str = "none",
+) -> TReal:
+    """gelu(Tensor self, *, str approximate="none") -> Tensor"""
+    return op20.Gelu(self, approximate=approximate)
+
+
+@onnx_impl(aten.group_norm.default, trace_only=True, opset_introduced=21)
+def aten_group_norm(
+    input: TFloat,
+    num_groups: int,
+    weight: TFloat | None = None,
+    bias: TFloat | None = None,
+    eps: float = 1e-05,
+    cudnn_enabled: bool = True,
+) -> TFloat:
+    """group_norm(Tensor input, int num_groups, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enabled=True) -> Tensor"""
+
+    c = op21.Shape(input, start=1, end=2)
+    if weight is None:
+        weight = op21.ConstantOfShape(c, value=ir.tensor([1.0], dtype=input.dtype))
+    if bias is None:
+        bias = op21.ConstantOfShape(c, value=ir.tensor([0.0], dtype=input.dtype))
+    return op21.GroupNormalization(
+        input, weight, bias, epsilon=eps, num_groups=num_groups
+    )
+
+
+@onnx_impl(aten.rms_norm.default, trace_only=True, opset_introduced=23)
+def aten_rms_norm(
+    input: TFloat,
+    normalized_shape: Sequence[int],
+    weight: TFloat | None = None,
+    eps: float | None = None,
+) -> TFloat:
+    """rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor"""
+
+    # Default eps value if not provided
+    if eps is None:
+        eps = torch.finfo(torch.float).eps  # Observed from decomp
+
+    # Calculate axis: the first normalization dimension
+    # For normalized_shape with D dimensions, normalize over last D dimensions
+    # Since ONNX RMSNormalization supports negative axis values, we use -len(normalized_shape)
+    # which correctly maps to the first axis of the normalized dimensions
+    normalized_dims = len(normalized_shape)
+    axis = -normalized_dims
+
+    # Create weight tensor if not provided
+    if weight is None:
+        weight = op23.ConstantOfShape(
+            op23.Shape(input), value=ir.tensor([1], dtype=input.dtype)
+        )
+
+    return op23.RMSNormalization(input, weight, axis=axis, epsilon=eps)
+
+
+@onnx_impl(
+    aten.scaled_dot_product_attention.default, trace_only=True, opset_introduced=23
+)
+def aten_scaled_dot_product_attention_23(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: TFloat | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float | None = None,
+    enable_gqa: bool = False,
+) -> TFloat:
+    """scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor
+
+    Reference:
+        1. https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+        2. https://onnx.ai/onnx/operators/onnx__Attention.html
+
+    Attempts to convert SDPA to Attention onnx op and fallbacks to an onnx graph equivalent to the following PyTorch code::
+        scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale
+        attn_mask = (
+            torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+            if is_causal
+            else attn_mask
+        )
+        attn_mask = (
+            attn_mask.masked_fill(not attn_mask, -float("inf"))
+            if attn_mask.dtype == torch.bool
+            else attn_mask
+        )
+        attn_weight = torch.softmax(
+            (Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1
+        )
+        attn_weight = torch.dropout(attn_weight, dropout_p)
+        return attn_weight @ V
+
+    where Q, K, V are the query, key, and value tensors, respectively.
+    L is the target sequence length, S is the source sequence length, and E is the embedding size.
+    """
+    if is_causal and attn_mask is not None:
+        raise AssertionError("is_causal and attn_mask cannot be set at the same time")
+    if not (len(query.shape) == 4 and len(key.shape) == 4 and len(value.shape) == 4):
+        raise AssertionError("only 4D query, key, and value are supported")
+
+    # Attention onnx op can only handle non-training scenarios where dropout is disabled.
+    if dropout_p == 0:
+        if enable_gqa:
+            if not (
+                query.shape[1] > key.shape[1] == value.shape[1]
+                and query.shape[1] % key.shape[1] == 0
+            ):
+                raise AssertionError(
+                    "SDPA (GQA or MQA) requires q_num_heads > kv_num_heads & "
+                    "q_num_heads % kv_num_heads == 0"
+                )
+        else:
+            if not (query.shape[1] == key.shape[1] == value.shape[1]):
+                raise AssertionError("SDPA (MHA) requires q_num_heads = kv_num_heads")
+
+        # NOTE: num_heads attributes (q_num_heads/kv_num_heads) should not be specified for 4D.
+        # They are not populated with 4D inputs because this information directly comes from input shapes:
+        # `q_num_heads=query.shape[1]` and `kv_num_heads=key.shape[1]`.
+        # This dimension is usually static but it could not be dynamic if also given as an attribute.
+        # num_heads attributes are needed for 3D attention inputs:
+        # (shape: [B, S, N*H]), 4D shape is ([B, N, S, H]).
+
+        Y, _, _, _ = op23.Attention(
+            query,
+            key,
+            value,
+            attn_mask=attn_mask,
+            scale=scale,
+            is_causal=is_causal,
+        )
+        return Y
+
+    if scale is None:
+        scale = _attention_scale(query, op23)
+    scale = op23.CastLike(scale, query)
+
+    if is_causal:
+        attn_mask = _causal_attention_mask(query, key, op23)
+
+    if enable_gqa:
+        key, value = _attention_repeat_kv_for_group_query(query, key, value, op23)
+
+    if attn_mask is None:
+        return _aten_scaled_dot_product_attention_no_mask_onnx(
+            query, key, value, scale, dropout_p, op23
+        )
+
+    return _aten_scaled_dot_product_attention_float_mask_onnx(
+        query, key, value, attn_mask, scale, dropout_p, op23
+    )
+
+
+def _attention_repeat_kv_for_group_query(
+    query: TFloat, key: TFloat, value: TFloat, op: Opset
+) -> tuple[TFloat, TFloat]:
+    """Expand key and value for group query attention.
+
+    repeat_interleave is applied on key and value to match the number of heads in query.
+
+    Args:
+        query: Tensor of shape [B, q_num_heads, q_S, E]
+        key: Tensor of shape [B, k_num_heads, kv_S, E]
+        value: Tensor of shape [B, v_num_heads, kv_S, E]
+
+    Returns:
+        Tuple of (expanded_key, expanded_value) where:
+            - expanded_key: Tensor of shape [B, q_num_heads, kv_S, E]
+            - expanded_value: Tensor of shape [B, q_num_heads, kv_S, E]
+    """
+
+    if not (
+        query.shape[1] > key.shape[1] == value.shape[1]
+        and query.shape[1] % key.shape[1] == 0
+    ):
+        raise AssertionError(
+            "SDPA (GQA or MQA) requires q_num_heads > kv_num_heads & "
+            "q_num_heads % kv_num_heads == 0"
+        )
+
+    # NOTE: QKV are expected to be 4D tensors
+
+    batch_size = op.Shape(query, start=0, end=1)  # [B]
+    q_num_heads = op.Shape(query, start=1, end=2)  # [Hq]
+    kv_num_heads = op.Shape(key, start=1, end=2)  # [Hk]
+    qk_head_size = op.Shape(key, start=3, end=4)  # [Dk]
+    v_head_size = op.Shape(value, start=3, end=4)  # [Dv]
+    new_kv_seq_len = op.Shape(key, start=2, end=3)  # [T]
+
+    interleave_dim = op.Div(q_num_heads, kv_num_heads)  # Hq / Hk
+    two = op.Constant(value_int=2)
+    k_unsqueezed = op.Unsqueeze(key, two)  # [B, Hk, 1, T, Dk]
+    v_unsqueezed = op.Unsqueeze(value, two)  # [B, Hv, 1, T, Dv]
+
+    k_expand_shape = op.Concat(
+        batch_size, kv_num_heads, interleave_dim, new_kv_seq_len, qk_head_size, axis=0
+    )
+    k_expand = op.Expand(k_unsqueezed, k_expand_shape)
+    v_expand_shape = op.Concat(
+        batch_size, kv_num_heads, interleave_dim, new_kv_seq_len, v_head_size, axis=0
+    )
+    v_expand = op.Expand(v_unsqueezed, v_expand_shape)
+
+    k_attention_shape = op.Concat(
+        batch_size, q_num_heads, new_kv_seq_len, qk_head_size, axis=0
+    )
+    v_attention_shape = op.Concat(
+        batch_size, q_num_heads, new_kv_seq_len, v_head_size, axis=0
+    )
+
+    expanded_key = op.Reshape(k_expand, k_attention_shape)
+    expanded_value = op.Reshape(v_expand, v_attention_shape)
+
+    return expanded_key, expanded_value
+
+
+def _attention_scale(query: TFloat, op: Opset) -> TFloat:
+    """Calculate the scale factor for the attention result.
+
+    Args:
+        query: Tensor of shape [..., L, E]
+
+    Returns:
+        Scalar scale factor := 1 / math.sqrt(query.size(-1))
+    """
+    q_shape = op.Shape(query)
+    q_last_dim = op.Gather(q_shape, op.Constant(value_ints=[-1]))
+    embedding_size = op.CastLike(q_last_dim, query)
+    one = op.Constant(value_float=1.0)
+    cast_one = op.CastLike(one, query)
+    scale = op.Div(cast_one, op.Sqrt(embedding_size))
+    return scale
+
+
+def _causal_attention_mask(query: TFloat, key: TFloat, op: Opset) -> TFloat:
+    """Create a causal mask for the given query and key tensors.
+
+    Equivalent to::
+        mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_mask = torch.zeros(L, S, dtype=torch.float)
+        attn_mask = attn_mask.masked_fill(not mask, -float("inf"))
+
+    Args:
+        query: Tensor of shape [..., L, E]
+        key: Tensor of shape [..., S, E]
+
+    Returns:
+        Tensor of shape [L, S]
+    """
+    q_shape = op.Shape(query)
+    k_shape = op.Shape(key)
+
+    target_length = op.Slice(
+        q_shape, op.Constant(value_ints=[-2]), op.Constant(value_ints=[-1])
+    )
+    source_length = op.Slice(
+        k_shape, op.Constant(value_ints=[-2]), op.Constant(value_ints=[-1])
+    )
+    # attn_mask = torch.ones(L, S) := {
+    size = op.Concat(target_length, source_length, axis=0)
+    attn_mask = op.Expand(op.Constant(value_float=1.0), size)
+    # }
+    attn_mask = op.Trilu(attn_mask, upper=0)
+    # The causal mask has 0s in the lower triangle and -inf in the upper triangle.
+    attn_mask = op.Where(
+        op.Equal(attn_mask, op.Constant(value_float=0.0)),
+        op.Constant(value_float=-float("inf")),
+        op.Constant(value_float=0.0),
+    )
+    attn_mask = op.CastLike(attn_mask, query)
+    return attn_mask
+
+
+def _aten_scaled_dot_product_attention_no_mask_onnx(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    scale: TFloat,
+    dropout_p: float,
+    op: Opset,
+) -> TFloat:
+    # Swap the last two axes of key
+    key_last_dim = op.Shape(key, start=-1)
+    key_second_last_dim = op.Shape(key, start=-2, end=-1)
+    key_first_dims = op.Shape(key, end=-2)
+    # Contract the dimensions that are not the last two so we can transpose
+    # with a static permutation.
+    key_squeezed_shape = op.Concat(
+        op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0
+    )
+    key_squeezed = op.Reshape(key, key_squeezed_shape)
+    key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1])
+    key_transposed_shape = op.Concat(
+        key_first_dims, key_last_dim, key_second_last_dim, axis=0
+    )
+    key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    query_scaled = op.Mul(query, op.Sqrt(scale))
+    key_transposed_scaled = op.Mul(
+        key_transposed, op.CastLike(op.Sqrt(scale), key_transposed)
+    )
+    attn_weight = op.Softmax(
+        op.MatMul(query_scaled, key_transposed_scaled),
+        axis=-1,
+    )
+    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    return op.MatMul(attn_weight, value)
+
+
+def _aten_scaled_dot_product_attention_float_mask_onnx(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: TFloat,
+    scale: TFloat,
+    dropout_p: float,
+    op: Opset,
+) -> TFloat:
+    # Swap the last two axes of key
+    key_last_dim = op.Shape(key, start=-1)
+    key_second_last_dim = op.Shape(key, start=-2, end=-1)
+    key_first_dims = op.Shape(key, end=-2)
+    # Contract the dimensions that are not the last two so we can transpose
+    # with a static permutation.
+    key_squeezed_shape = op.Concat(
+        op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0
+    )
+    key_squeezed = op.Reshape(key, key_squeezed_shape)
+    key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1])
+    key_transposed_shape = op.Concat(
+        key_first_dims, key_last_dim, key_second_last_dim, axis=0
+    )
+    key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    query_scaled = op.Mul(query, op.Sqrt(scale))
+    key_transposed_scaled = op.Mul(key_transposed, op.Sqrt(scale))
+    attn_weight = op.Softmax(
+        op.Add(op.MatMul(query_scaled, key_transposed_scaled), attn_mask),
+        axis=-1,
+    )
+    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    return op.MatMul(attn_weight, value)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symbolic.py

@@ -0,0 +1,150 @@
+"""Implementation for higher-order operators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from onnxscript.ir import convenience as ir_convenience
+
+import torch
+from torch.onnx._internal._lazy_import import onnx_ir as ir
+from torch.onnx._internal.exporter import _core
+from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl
+from torch.onnx.ops import _symbolic_impl
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+def _call_symbolic_op(
+    op_type: str,
+    domain: str,
+    args: Sequence[ir.Value | None],
+    kwargs: dict[str, int | float | str | bool | list[int] | list[float] | list[str]],
+    dtypes: Sequence[int],
+    version: int | None,
+    metadata_props: dict[str, str] | None,
+) -> Sequence[ir.Value]:
+    """Call an operator with the given arguments and keyword arguments.
+
+    Arguments are always inputs, while keyword arguments are attributes.
+    """
+    # This is a wrapper around the IR node creation that hooks into the _builder.OpRecorder
+    # tracer so that all nodes created are recorded the same way as if we were to use
+    # onnxscript ops directly.
+
+    if _core.current_tracer is None:
+        raise AssertionError("current_tracer must be non-None")
+    tracer = _core.current_tracer
+
+    inputs = list(args)
+
+    # If final inputs are None, strip them from the node inputs
+    for input in reversed(inputs):
+        if input is not None:
+            break
+        inputs.pop()
+
+    # Construct and filter out None attributes
+    attributes = [
+        attr
+        for attr in ir_convenience.convert_attributes(kwargs)  # type: ignore[arg-type]
+        if attr.value is not None  # type: ignore[union-attr]
+    ]
+    tracer.nodes.append(
+        node := ir.Node(
+            domain,
+            op_type,
+            inputs=inputs,
+            attributes=attributes,
+            num_outputs=len(dtypes),
+            version=version,
+            metadata_props=metadata_props,
+        )
+    )
+    # Set the dtypes for the outputs. We set them here because the graph builder
+    # Uses PyTorch types which are sometimes inaccurate when they are ONNX only
+    # types like float4e2m1.
+    for value, dtype in zip(node.outputs, dtypes):
+        value.dtype = ir.DataType(dtype)
+        # The shape is set by the graph builder. We don't need to set it here.
+    return node.outputs
+
+
+@onnx_impl(torch.ops.onnx_symbolic._symbolic.default, no_compile=True)
+def onnx_symbolic_symbolic(
+    inputs: Sequence[ir.Value | None],
+    op_type: str,
+    onnx_dtype: int,
+    *,
+    shape: Sequence[int | ir.Value],
+    attr_keys: Sequence[str],
+    attr_types: Sequence[str],
+    attr_pos: Sequence[tuple[int, int]],
+    attr_ints: Sequence[int],
+    attr_floats: Sequence[float],
+    attr_strs: Sequence[str],
+    metadata_props_keys: Sequence[str] = (),
+    metadata_props_values: Sequence[str] = (),
+    domain: str = "",
+    version: int | None = None,
+) -> ir.Value:
+    del shape  # Unused. The shapes are set by the graph builder
+    encoded = _symbolic_impl.EncodedAttrs(
+        attr_keys=list(attr_keys),
+        attr_types=list(attr_types),
+        attr_pos=list(attr_pos),
+        attr_ints=list(attr_ints),
+        attr_floats=list(attr_floats),
+        attr_strs=list(attr_strs),
+    )
+    attrs = encoded.to_dict()
+    return _call_symbolic_op(
+        op_type,
+        domain,
+        inputs,
+        attrs,
+        dtypes=[onnx_dtype],
+        version=version,
+        metadata_props=dict(zip(metadata_props_keys, metadata_props_values)),
+    )[0]
+
+
+@onnx_impl(torch.ops.onnx_symbolic._symbolic_multi_out.default, no_compile=True)
+def onnx_symbolic_symbolic_multi_out(
+    inputs: Sequence[ir.Value | None],
+    op_type: str,
+    onnx_dtypes: Sequence[int],
+    *,
+    shapes: Sequence[Sequence[int | ir.Value]],
+    attr_keys: Sequence[str],
+    attr_types: Sequence[str],
+    attr_pos: Sequence[tuple[int, int]],
+    attr_ints: Sequence[int],
+    attr_floats: Sequence[float],
+    attr_strs: Sequence[str],
+    metadata_props_keys: Sequence[str] = (),
+    metadata_props_values: Sequence[str] = (),
+    domain: str = "",
+    version: int | None = None,
+) -> Sequence[ir.Value]:
+    del shapes  # Unused. The shapes are set by the graph builder
+    encoded = _symbolic_impl.EncodedAttrs(
+        attr_keys=list(attr_keys),
+        attr_types=list(attr_types),
+        attr_pos=list(attr_pos),
+        attr_ints=list(attr_ints),
+        attr_floats=list(attr_floats),
+        attr_strs=list(attr_strs),
+    )
+    attrs = encoded.to_dict()
+    return _call_symbolic_op(
+        op_type,
+        domain,
+        inputs,
+        attrs,
+        dtypes=onnx_dtypes,
+        version=version,
+        metadata_props=dict(zip(metadata_props_keys, metadata_props_values)),
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symops.py

@@ -0,0 +1,64 @@
+"""Implementation for torch.sym* ops."""
+
+# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index"
+# pyrefly: ignore-errors
+# ruff: noqa: TCH001,TCH002,TC003
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from onnxscript.onnx_opset import opset18 as op
+
+import torch
+from torch.onnx._internal.exporter._torchlib._tensor_typing import (
+    BOOL,
+    FLOAT,
+    IntType,
+    TensorType,
+    TTensor,
+)
+from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl
+
+
+@onnx_impl(torch.sym_float, trace_only=True)
+def sym_float(self: TensorType) -> FLOAT:
+    """sym_float(SymInt self) -> SymFloat"""
+    return op.Cast(self, to=FLOAT.dtype)
+
+
+@onnx_impl(torch.sym_max, trace_only=True)
+def sym_max(x: IntType, y: IntType) -> IntType:
+    """sym_max(SymInt x, SymInt y) -> SymInt"""
+    return op.Max(x, y)
+
+
+@onnx_impl(torch.sym_min, trace_only=True)
+def sym_min(x: IntType, y: IntType) -> IntType:
+    """sym_min(SymInt x, SymInt y) -> SymInt"""
+    return op.Min(x, y)
+
+
+@onnx_impl(torch.sym_not, trace_only=True)
+def sym_not(self: BOOL) -> BOOL:
+    """sym_not(SymBool self) -> SymBool"""
+    return op.Not(self)
+
+
+@onnx_impl(torch.sym_sum, trace_only=True)
+def sym_sum(args: Sequence[IntType]) -> IntType:
+    """sym_sum(SymInt[] args) -> SymInt"""
+    if len(args) == 0:
+        return op.Constant(value_int=0)
+    if len(args) == 1:
+        return args[0]
+    result = op.Add(args[0], args[1])
+    for i in range(2, len(args)):
+        result = op.Add(result, args[i])
+    return result
+
+
+@onnx_impl(torch.sym_ite, trace_only=True)
+def sym_ite(b: BOOL, t: TTensor, f: TTensor) -> TTensor:
+    """sym_ite(SymBool b, Tensor t, Tensor f) -> Tensor"""
+    return op.Where(b, t, f)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_type_casting.py

@@ -0,0 +1,34 @@
+import numpy as np
+
+import torch
+
+
+def unpack_float4x2_as_uint8(tensor: torch.Tensor) -> np.ndarray:
+    """Convert a float4x2 tensor to unpacked uint8 np array."""
+    if tensor.dtype != torch.float4_e2m1fn_x2:
+        raise AssertionError(f"Expected float4_e2m1fn_x2, got {tensor.dtype}")
+    data = tensor.view(torch.uint8).numpy(force=True).flatten()
+    result_size = tensor.numel() * 2
+    result = np.empty([result_size], dtype=np.uint8)
+    array_low = data & np.uint8(0x0F)
+    array_high = data & np.uint8(0xF0)
+    array_high >>= np.uint8(4)
+    result[0::2] = array_low
+    result[1::2] = array_high
+    result.resize(get_float4_shape(tensor), refcheck=False)
+    return result
+
+
+def get_float4_shape(tensor: torch.Tensor) -> tuple[int, ...]:
+    """Get the shape of an unpacked float4 tensor.
+
+    The float4_e2m1fn_x2 type is a shell type described in
+    https://github.com/pytorch/pytorch/issues/146414.
+
+    the shell dtype is takes up 1 byte per element and semantically represents
+    two fp4 values packed into 1 byte. Semantically it represents (*tensor.shape[:-1], tensor.shape[-1]*2)
+    fp4 elements.
+    """
+    if tensor.dtype != torch.float4_e2m1fn_x2:
+        raise AssertionError(f"Expected float4_e2m1fn_x2, got {tensor.dtype}")
+    return (*tensor.shape[:-1], tensor.shape[-1] * 2)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_verification.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+
+__all__ = [
+    "VerificationInfo",
+    "verify_onnx_program",
+]
+
+import dataclasses
+import logging
+import math
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch.utils import _pytree
+
+
+if TYPE_CHECKING:
+    from torch.onnx._internal._lazy_import import onnx_ir as ir
+    from torch.onnx._internal.exporter import _onnx_program
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class VerificationInfo:
+    """Verification information for a value in the ONNX program.
+
+    This class contains the maximum absolute difference, maximum relative difference,
+    and histograms of absolute and relative differences between the expected and actual
+    values. It also includes the expected and actual data types.
+
+    The histograms are represented as tuples of tensors, where the first tensor is the
+    histogram counts and the second tensor is the bin edges.
+
+    Attributes:
+        name: The name of the value (output or intermediate).
+        max_abs_diff: The maximum absolute difference between the expected and actual values.
+        max_rel_diff: The maximum relative difference between the expected and actual values.
+        abs_diff_hist: A tuple of tensors representing the histogram of absolute differences.
+            The first tensor is the histogram counts and the second tensor is the bin edges.
+        rel_diff_hist: A tuple of tensors representing the histogram of relative differences.
+            The first tensor is the histogram counts and the second tensor is the bin edges.
+        expected_dtype: The data type of the expected value.
+        actual_dtype: The data type of the actual value.
+    """
+
+    name: str
+    max_abs_diff: float
+    max_rel_diff: float
+    abs_diff_hist: tuple[torch.Tensor, torch.Tensor]
+    rel_diff_hist: tuple[torch.Tensor, torch.Tensor]
+    expected_dtype: torch.dtype
+    actual_dtype: torch.dtype
+    # NOTE: We don't need to include shape because the expected shape is already known
+    # and checked by the runtime
+
+    @classmethod
+    def from_tensors(
+        cls,
+        name: str,
+        expected: torch.Tensor | float | int | bool,
+        actual: torch.Tensor | float | int | bool,
+    ) -> VerificationInfo:
+        """Create a VerificationInfo object from two tensors.
+
+        Args:
+            name: The name of the value.
+            expected: The expected tensor.
+            actual: The actual tensor.
+
+        Returns:
+            VerificationInfo: The VerificationInfo object.
+        """
+        if not isinstance(expected, torch.Tensor):
+            expected = torch.tensor(expected)
+        if not isinstance(actual, torch.Tensor):
+            actual = torch.tensor(actual)
+
+        max_abs_diff, max_rel_diff, abs_diff, rel_diff = _compare_tensors(
+            expected, actual
+        )
+        bins = torch.tensor(
+            [0.0, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1.0, 10, 1000000],
+            dtype=torch.float,
+        )
+        abs_diff_hist = torch.histogram(abs_diff.float(), bins=bins)
+        rel_diff_hist = torch.histogram(rel_diff.float(), bins=bins)
+        return cls(
+            name=name,
+            max_abs_diff=max_abs_diff,
+            max_rel_diff=max_rel_diff,
+            abs_diff_hist=abs_diff_hist,
+            rel_diff_hist=rel_diff_hist,
+            expected_dtype=expected.dtype,
+            actual_dtype=actual.dtype,
+        )
+
+    def asdict(self) -> dict[str, Any]:
+        """Convert the VerificationInfo object to a dictionary.
+
+        Returns:
+            A dictionary representation of the VerificationInfo object.
+        """
+        return {
+            "name": self.name,
+            "max_abs_diff": self.max_abs_diff,
+            "max_rel_diff": self.max_rel_diff,
+            "abs_diff_hist": [
+                self.abs_diff_hist[0].tolist(),
+                self.abs_diff_hist[1].tolist(),
+            ],
+            "rel_diff_hist": [
+                self.rel_diff_hist[0].tolist(),
+                self.rel_diff_hist[1].tolist(),
+            ],
+            "expected_dtype": str(self.expected_dtype),
+            "actual_dtype": str(self.actual_dtype),
+        }
+
+
+def _compare_tensors(
+    expected: torch.Tensor,
+    actual: torch.Tensor,
+) -> tuple[float, float, torch.Tensor, torch.Tensor]:
+    # Move tensors to the same device
+    expected = expected.detach().cpu()
+    actual = actual.detach().cpu()
+    if expected.numel() == 0 or actual.numel() == 0:
+        return math.inf, math.inf, torch.tensor(math.inf), torch.tensor(math.inf)
+    if expected.dtype == torch.bool:
+        expected = expected.to(torch.float32)
+        actual = actual.to(torch.float32)
+    if torch.is_complex(expected):
+        expected = torch.view_as_real(expected)
+    abs_diff = torch.abs(expected - actual)
+    eps = 1e-7
+    normalizer = torch.abs(expected) + eps
+    rel_diff = abs_diff / normalizer
+
+    max_absolute_difference = abs_diff.max().item()
+    max_relative_difference = rel_diff.max().item()
+
+    return max_absolute_difference, max_relative_difference, abs_diff, rel_diff
+
+
+def verify_onnx_program(
+    onnx_program: _onnx_program.ONNXProgram,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    compare_intermediates: bool = False,
+) -> list[VerificationInfo]:
+    """Verify the ONNX model by comparing the values with the expected values from ExportedProgram.
+
+    Args:
+        onnx_program: The ONNX program to verify.
+        args: The input arguments for the model.
+        kwargs: The keyword arguments for the model.
+        compare_intermediates: Whether to verify intermediate values. This is going
+            to take longer time, so it is disabled by default.
+
+    Returns:
+        VerificationInfo objects containing the verification information for each value.
+    """
+    exported_program = onnx_program.exported_program
+    if exported_program is None:
+        raise ValueError(
+            "The ONNX program does not contain an exported_program. "
+            "Please provide an exported_program to verify the ONNX program."
+        )
+    if args is None and kwargs is None:
+        # User did not provide example inputs, use the default example inputs
+        if exported_program.example_inputs is None:
+            raise ValueError(
+                "No example inputs provided and the exported_program does not contain example inputs. "
+                "Please provide arguments to verify the ONNX program."
+            )
+        args, kwargs = exported_program.example_inputs
+    if args is None:
+        args = ()
+    if kwargs is None:
+        kwargs = {}
+
+    # Flatten args for ONNX program and the VerificationInterpreter
+    flat_args, _ = exported_program._get_flat_args_with_check(args, kwargs)
+
+    if not compare_intermediates:
+        # Compare the output values
+        torch_outputs, _ = _pytree.tree_flatten(
+            exported_program.module()(*args, **kwargs)
+        )
+        onnx_outputs = onnx_program(*flat_args)
+        results = []
+        for torch_output, onnx_output, output_val in zip(
+            torch_outputs, onnx_outputs, onnx_program.model.graph.outputs
+        ):
+            results.append(
+                VerificationInfo.from_tensors(
+                    name=str(output_val.name),
+                    expected=torch_output,
+                    actual=onnx_output,
+                )
+            )
+        return results
+
+    # Use the _VerificationInterpreter to get the intermediate values
+    # By design the output values are included too
+    interpreter = _VerificationInterpreter(onnx_program)
+    interpreter.run(*flat_args)
+
+    return interpreter.verification_infos
+
+
+def _create_value_mapping(graph: ir.Graph) -> dict[str, ir.Value]:
+    """Return a dictionary mapping names to values in the graph.
+
+    The mapping does not include values from subgraphs.
+
+    Args:
+        graph: The graph to extract the mapping from.
+
+    Returns:
+        A dictionary mapping names to values.
+    """
+    values: dict[str, ir.Value] = {}
+    values.update(graph.initializers)
+    # The names of the values can be None or "", which we need to exclude
+    for input in graph.inputs:
+        if not input.name:
+            continue
+        values[input.name] = input
+    for node in graph:
+        for value in node.outputs:
+            if not value.name:
+                continue
+            values[value.name] = value
+    return values
+
+
+class _VerificationInterpreter(torch.fx.Interpreter):
+    """Interpreter for verifying converted ONNX model accuracy by comparing intermediate values.
+
+    To compare models, first initialize the interpreter with an ONNX program.
+    Then, call the :meth:`run` method with the input arguments to execute the model.
+    The :meth:`run` method will execute the model and populate the
+    :attr:`verification_infos` attribute with the verification information for each value.
+
+    ::
+        onnx_program = torch.onnx.export(model, args, dynamo=True)
+        interpreter = _VerificationInterpreter(onnx_program)
+        interpreter.run(*args)
+        verification_infos = interpreter.verification_infos
+        for info in verification_infos:
+            print("value name:", info.name, info)
+
+    The verification information includes the maximum absolute difference, maximum relative
+    difference, and histograms of absolute and relative differences between the expected
+    and actual values. See :class:`VerificationInfo` for more details.
+
+    Attributes:
+        verification_infos: A list of verification information for each value.
+            It is populated when the `run` method is called.
+    """
+
+    def __init__(self, onnx_program: torch.onnx.ONNXProgram) -> None:
+        """Initialize the _VerificationInterpreter with an ONNX program.
+
+        Args:
+            onnx_program: The ONNX program to verify.
+        """
+        if onnx_program.exported_program is None:
+            raise ValueError(
+                "The ONNX program does not contain an exported_program. "
+                "Please provide an exported_program to verify the ONNX program."
+            )
+        super().__init__(onnx_program.exported_program.module())
+        self._onnx_program = onnx_program
+        self._onnx_values = _create_value_mapping(onnx_program.model.graph)
+        self._args: tuple[Any, ...] = ()
+        self.verification_infos: list[VerificationInfo] = []
+
+    def run(
+        self,
+        *args: Any,
+        initial_env: dict[torch.fx.Node, Any] | None = None,
+        enable_io_processing: bool = True,
+    ) -> Any:
+        """Run the interpreter with the given input arguments.
+
+        This method executes the model and populates the :attr:`verification_infos` attribute
+        with the verification information for each value.
+
+        Args:
+            args: The input arguments for the model.
+            initial_env: The initial environment for the interpreter.
+            enable_io_processing: Whether to enable IO processing.
+
+        Returns:
+            Any: The result of executing the model.
+        """
+        self.verification_infos = []
+        self._args = args
+        return super().run(
+            *args,
+            initial_env=initial_env,
+            enable_io_processing=enable_io_processing,
+        )
+
+    def run_node(self, n: torch.fx.Node) -> Any:
+        result = super().run_node(n)
+        if n.op != "call_function":
+            return result
+        node_name = n.name
+        if node_name not in self._onnx_values:
+            return result
+        try:
+            (onnx_result,) = self._onnx_program.compute_values([node_name], self._args)
+        except Exception:
+            logger.warning(
+                "Failed to compute value for node %s", node_name, exc_info=True
+            )
+            return result
+        info = VerificationInfo.from_tensors(
+            name=node_name,
+            expected=result,
+            actual=onnx_result,
+        )
+        self.verification_infos.append(info)
+        if info.max_abs_diff > 0.01 or info.max_rel_diff > 0.1:
+            logger.warning(
+                "Verification info for node %s: max_abs_diff: %s, max_rel_diff: %s",
+                node_name,
+                info.max_abs_diff,
+                info.max_rel_diff,
+            )
+        else:
+            logger.info(
+                "Verification info for node %s: max_abs_diff: %s, max_rel_diff: %s",
+                node_name,
+                info.max_abs_diff,
+                info.max_rel_diff,
+            )
+        return result

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/fx/_pass.py

@@ -0,0 +1,237 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import abc
+import contextlib
+import dataclasses
+import difflib
+import io
+import sys
+from typing import Any, TYPE_CHECKING
+
+import torch
+import torch.fx
+from torch._subclasses.fake_tensor import unset_fake_temporarily
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch._subclasses import fake_tensor
+
+
+@dataclasses.dataclass
+class PackageInfo:
+    package_name: str
+    version: str | None
+    commit_hash: str | None
+
+    def to_onnx_domain_string(self) -> str:
+        return ".".join(
+            filter(None, ("pkg", self.package_name, self.version, self.commit_hash))
+        )
+
+    @classmethod
+    def from_python_class(cls, python_class_name: type | str) -> PackageInfo:
+        if isinstance(python_class_name, type):
+            python_class_name = python_class_name.__module__
+        package_name = python_class_name.split(".")[0]
+        package = __import__(package_name)
+        version = getattr(package, "__version__", None)
+        # TODO: Figure out how to retrieve commit hash.
+        commit_hash = None
+        return cls(package_name, version, commit_hash)
+
+
+@dataclasses.dataclass
+class GraphModuleOnnxMeta:
+    package_info: PackageInfo
+
+
+@contextlib.contextmanager
+def _patch_difflib_sequence_matcher_init():
+    """Context patching `difflib.SequenceMatcher` for fx readable graph.
+
+    Under this context, the `autojunk` argument of `difflib.SequenceMatcher` will always
+    be considered as `False`. This is to prevent `difflib.SequenceMatcher` recognizing
+    stacktrace messages in fx readable graph as junk, as these messages tend to be long (>200)
+    and repeat multiple times, which falls under the junk filter criteria.
+
+    `difflib.SequenceMatcher` is used underneath by all sorts of diffing functions
+    in `difflib`, including `difflib.unified_diff`, `difflib.ndiff`, `difflib.context_diff`.
+    Unfortunately, there is no way to pass `autojunk` argument to these functions, and
+    they all default to `True`. This context patching will affect all of them.
+
+    `Reference: Automatic junk heuristic <https://docs.python.org/3/library/difflib.html>`_
+    """
+    original_init = difflib.SequenceMatcher.__init__
+
+    def patched_init(self, isjunk=None, a="", b="", autojunk=True) -> None:
+        original_init(self, isjunk, a, b, autojunk=False)
+
+    difflib.SequenceMatcher.__init__ = patched_init  # type: ignore[assignment]
+    try:
+        yield
+    finally:
+        difflib.SequenceMatcher.__init__ = original_init  # type: ignore[assignment]
+
+
+def _unified_diff(a: str, b: str) -> str:
+    """Return a string containing the unified diff of two strings.
+
+    This function calls a patched version of `difflib.unified_diff` with `autojunk` set
+    to `False` for `difflib.SequenceMatcher` class. More details can be found in
+    `_patch_difflib_sequence_matcher_init` function.
+
+    Args:
+        a: The first string.
+        b: The second string.
+
+    Returns:
+        The unified diff of the two strings. If there is no diff, return "<no diff>".
+
+    Example::
+
+        >>> a = '''class GraphModule(torch.nn.Module):
+        ...     def forward(self, input_ids : torch.Tensor, attention_mask : torch.Tensor):
+        ...         # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1])
+        ...         view = input_ids.view(-1, 3);  input_ids = None
+        ... '''
+        >>> b = '''class <lambda>(torch.nn.Module):
+        ...     def forward(self, input_ids: i64[1, 3], attention_mask: i64[1, 3]):
+        ...         # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1])
+        ...         view: i64[1, 3] = torch.ops.aten.view.default(input_ids, [-1, 3]);  input_ids = None
+        ... '''
+        >>> print(_unified_diff(a, b))
+        ---
+        +++
+        @@ -1,4 +1,4 @@
+        -class GraphModule(torch.nn.Module):
+        -    def forward(self, input_ids : torch.Tensor, attention_mask : torch.Tensor):
+        +class <lambda>(torch.nn.Module):
+        +    def forward(self, input_ids: i64[1, 3], attention_mask: i64[1, 3]):
+                # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1])
+        -        view = input_ids.view(-1, 3);  input_ids = None
+        +        view: i64[1, 3] = torch.ops.aten.view.default(input_ids, [-1, 3]);  input_ids = None
+    """
+
+    a_list = a.splitlines(keepends=True)
+    b_list = b.splitlines(keepends=True)
+
+    with _patch_difflib_sequence_matcher_init():
+        # Set `n` to `sys.maxsize` to show entire graph when there is a diff.
+        diff = "".join(difflib.unified_diff(a_list, b_list, n=sys.maxsize))
+
+    if not diff:
+        return "<no diff>"
+    return diff
+
+
+def _transform_diagnose_call_message_formatter(
+    run: Callable,
+    self: Transform,
+    *args: Any,
+    **kwargs: Any,
+) -> str:
+    return f"Running {self.__class__.__name__} pass. "
+
+
+def maybe_fx_graph_tabular(graph: torch.fx.Graph) -> str | None:
+    """Return the Graph nodes in tabular format. Equivalent to stdout of `graph.print_tabular()`.
+    If `tabulate` is not installed, return `None`.
+
+    Args:
+        graph: The Graph to print.
+
+    Returns:
+        The Graph printed in a tabular format. None if `tabulate` is not installed.
+    """
+    f = io.StringIO()
+    with contextlib.redirect_stdout(f):
+        try:
+            graph.print_tabular()
+        except ImportError:
+            return None
+    return f.getvalue()
+
+
+class Transform(abc.ABC):
+    """Base class for FX graph transformations to be used by FX-ONNX exporter.
+
+    Similar to `FX Interpreter <https://pytorch.org/docs/stable/fx.html#torch.fx.Interpreter>`_,
+    specializations of this class execute the FX graph Node-by-Node.
+    Methods in the `Transform` class can be overridden to customize the behavior of the model.
+    This pattern can be useful for many things, including writing code transformations as well as analysis passes.
+
+    The following methods can be overridden::
+
+        _run()
+            +-- run_node()
+                +-- placeholder()
+                +-- get_attr()
+                +-- call_function()
+                +-- call_method()
+                +-- call_module()
+                +-- output()
+
+    One important aspect to note is that if the transformation modifies the model input and/or output signature,
+    (e.g. additional inputs/outputs are added to the model), :class:`InputAdaptStep` and/or :class:`OutputAdaptStep`
+    are needed to reconcile :attr:`ONNXProgram.model_proto`.
+    That is, the model signature and the model representation must match.
+
+    TODO(bowbao): Add more overridable methods in call hierarchy
+    TODO(bowbao): Create an example once more overridable methods are added.
+    """
+
+    module: torch.fx.GraphModule
+    """The module to be transformed."""
+
+    fake_mode: fake_tensor.FakeTensorMode | None
+    """The existing fake mode detected from `self.module`."""
+
+    def __init__(
+        self,
+        module: torch.fx.GraphModule,
+    ) -> None:
+        """Initialize the transform.
+
+        Args:
+            module: The module to be transformed.
+        """
+        self.module = module
+        self.fake_mode = self._detect_fake_mode()
+
+    def _detect_fake_mode(self) -> fake_tensor.FakeTensorMode | None:
+        """Detect fake mode from the graph.
+
+        Scan through all nodes in graph and their meta['val'] to detect fake mode.
+        """
+        fake_tensors = [node.meta.get("val") for node in self.module.graph.nodes]
+        with unset_fake_temporarily():
+            return torch._dynamo.utils.detect_fake_mode(fake_tensors)
+
+    def _maybe_fakefy_args(
+        self, fake_mode: fake_tensor.FakeTensorMode | None, *args: Any
+    ) -> tuple[Any, ...]:
+        if fake_mode is None:
+            return args
+        # NB: This should hit the cache if tensors were fakefied before.
+        # E.g., when the fx graph is produced by Dynamo.
+        return tuple(
+            fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args
+        )
+
+    @abc.abstractmethod
+    def _run(self, *args, **kwargs) -> torch.fx.GraphModule: ...
+
+    def run(self, *args, **kwargs) -> torch.fx.GraphModule:
+        """Run the transform on `self.module`.
+
+        Note that this method may or may not mutate `self.module`, and the returned
+        `GraphModule` could be either `self.module` or a new `GraphModule`.
+
+        Args:
+            *args: Positional arguments for `self.module` to run.
+            **kwargs: Keyword arguments for `self.module` to run.
+        """
+        return self._run(*args, **kwargs)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/fx/passes/__init__.py

@@ -0,0 +1,6 @@
+from .type_promotion import InsertTypePromotion
+
+
+__all__ = [
+    "InsertTypePromotion",
+]

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/fx/type_utils.py

@@ -0,0 +1,38 @@
+# mypy: allow-untyped-defs
+"""Utilities for converting and operating on ONNX and torch types."""
+
+from __future__ import annotations
+
+from typing import Any
+from typing_extensions import TypeIs
+
+import torch
+
+
+def is_torch_symbolic_type(
+    value: Any,
+) -> TypeIs[torch.SymBool | torch.SymInt | torch.SymFloat]:
+    return isinstance(value, (torch.SymBool, torch.SymInt, torch.SymFloat))
+
+
+def from_scalar_type_to_torch_dtype(scalar_type: type) -> torch.dtype | None:
+    return _SCALAR_TYPE_TO_TORCH_DTYPE.get(scalar_type)
+
+
+_PYTHON_TYPE_TO_TORCH_DTYPE = {
+    bool: torch.bool,
+    int: torch.int64,
+    float: torch.float32,
+    complex: torch.complex64,
+}
+
+_SYM_TYPE_TO_TORCH_DTYPE = {
+    torch.SymInt: torch.int64,
+    torch.SymFloat: torch.float32,
+    torch.SymBool: torch.bool,
+}
+
+_SCALAR_TYPE_TO_TORCH_DTYPE: dict[type, torch.dtype] = {
+    **_PYTHON_TYPE_TO_TORCH_DTYPE,
+    **_SYM_TYPE_TO_TORCH_DTYPE,  # type: ignore[dict-item]
+}

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/_experimental.py

@@ -0,0 +1,27 @@
+"""Experimental classes and functions used by ONNX export."""
+
+import dataclasses
+from collections.abc import Mapping, Sequence
+
+import torch
+import torch._C._onnx as _C_onnx
+
+
+@dataclasses.dataclass
+class ExportOptions:
+    """Arguments used by :func:`torch.onnx.export`."""
+
+    # TODO(justinchuby): Deprecate and remove this class.
+
+    export_params: bool = True
+    verbose: bool = False
+    training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL
+    input_names: Sequence[str] | None = None
+    output_names: Sequence[str] | None = None
+    operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX
+    opset_version: int | None = None
+    do_constant_folding: bool = True
+    dynamic_axes: Mapping[str, Mapping[int, str] | Sequence[int]] | None = None
+    keep_initializers_as_inputs: bool | None = None
+    custom_opsets: Mapping[str, int] | None = None
+    export_modules_as_functions: bool | set[type[torch.nn.Module]] = False

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/_globals.py

@@ -0,0 +1,81 @@
+"""Globals used internally by the ONNX exporter.
+
+Do not use this module outside of `torch.onnx` and its tests.
+
+Be very judicious when adding any new global variables. Do not create new global
+variables unless they are absolutely necessary.
+"""
+
+import torch._C._onnx as _C_onnx
+
+# This module should only depend on _constants and nothing else in torch.onnx to keep
+# dependency direction clean.
+from torch.onnx import _constants
+
+
+class _InternalGlobals:
+    """Globals used internally by ONNX exporter.
+
+    NOTE: Be very judicious when adding any new variables. Do not create new
+    global variables unless they are absolutely necessary.
+    """
+
+    def __init__(self) -> None:
+        self._export_onnx_opset_version = _constants.ONNX_DEFAULT_OPSET
+        self._training_mode: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL
+        self._in_onnx_export: bool = False
+        # Whether the user's model is training during export
+        self.export_training: bool = False
+        self.operator_export_type: _C_onnx.OperatorExportTypes = (
+            _C_onnx.OperatorExportTypes.ONNX
+        )
+        self.onnx_shape_inference: bool = True
+        self._autograd_inlining: bool = True
+
+    @property
+    def training_mode(self) -> _C_onnx.TrainingMode:
+        """The training mode for the exporter."""
+        return self._training_mode
+
+    @training_mode.setter
+    def training_mode(self, training_mode: _C_onnx.TrainingMode) -> None:
+        if not isinstance(training_mode, _C_onnx.TrainingMode):
+            raise TypeError(
+                "training_mode must be of type 'torch.onnx.TrainingMode'. This is "
+                "likely a bug in torch.onnx."
+            )
+        self._training_mode = training_mode
+
+    @property
+    def export_onnx_opset_version(self) -> int:
+        """Opset version used during export."""
+        return self._export_onnx_opset_version
+
+    @export_onnx_opset_version.setter
+    def export_onnx_opset_version(self, value: int) -> None:
+        self._export_onnx_opset_version = value
+
+    @property
+    def in_onnx_export(self) -> bool:
+        """Whether it is in the middle of ONNX export."""
+        return self._in_onnx_export
+
+    @in_onnx_export.setter
+    def in_onnx_export(self, value: bool) -> None:
+        if type(value) is not bool:
+            raise TypeError("in_onnx_export must be a boolean")
+        self._in_onnx_export = value
+
+    @property
+    def autograd_inlining(self) -> bool:
+        """Whether Autograd must be inlined."""
+        return self._autograd_inlining
+
+    @autograd_inlining.setter
+    def autograd_inlining(self, value: bool) -> None:
+        if type(value) is not bool:
+            raise TypeError("autograd_inlining must be a boolean")
+        self._autograd_inlining = value
+
+
+GLOBALS = _InternalGlobals()

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/_type_utils.py

@@ -0,0 +1,393 @@
+# mypy: allow-untyped-defs
+"""Utilities for converting and operating on ONNX, JIT and torch types."""
+
+from __future__ import annotations
+
+import enum
+import typing
+from typing import Literal
+
+import torch
+from torch._C import _onnx as _C_onnx
+from torch.onnx import errors
+
+
+if typing.TYPE_CHECKING:
+    # Hack to help mypy to recognize torch._C.Value
+    from torch import _C  # noqa: F401
+
+ScalarName = Literal[
+    "Byte",
+    "Char",
+    "Double",
+    "Float",
+    "Half",
+    "Int",
+    "Long",
+    "Short",
+    "Bool",
+    "ComplexHalf",
+    "ComplexFloat",
+    "ComplexDouble",
+    "QInt8",
+    "QUInt8",
+    "QInt32",
+    "BFloat16",
+    "Float8E5M2",
+    "Float8E4M3FN",
+    "Float8E5M2FNUZ",
+    "Float8E4M3FNUZ",
+    "Undefined",
+]
+
+TorchName = Literal[
+    "bool",
+    "uint8_t",
+    "int8_t",
+    "double",
+    "float",
+    "half",
+    "int",
+    "int64_t",
+    "int16_t",
+    "complex32",
+    "complex64",
+    "complex128",
+    "qint8",
+    "quint8",
+    "qint32",
+    "bfloat16",
+    "float8_e5m2",
+    "float8_e4m3fn",
+    "float8_e5m2fnuz",
+    "float8_e4m3fnuz",
+]
+
+
+class JitScalarType(enum.IntEnum):
+    """Scalar types defined in torch.
+
+    Use ``JitScalarType`` to convert from torch and JIT scalar types to ONNX scalar types.
+
+    Examples:
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
+        >>> # xdoctest: +IGNORE_WANT("win32 has different output")
+        >>> JitScalarType.from_value(torch.ones(1, 2)).onnx_type()
+        TensorProtoDataType.FLOAT
+
+        >>> JitScalarType.from_value(torch_c_value_with_type_float).onnx_type()
+        TensorProtoDataType.FLOAT
+
+        >>> JitScalarType.from_dtype(torch.get_default_dtype).onnx_type()
+        TensorProtoDataType.FLOAT
+
+    """
+
+    # Order defined in https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h
+    UINT8 = 0
+    INT8 = enum.auto()  # 1
+    INT16 = enum.auto()  # 2
+    INT = enum.auto()  # 3
+    INT64 = enum.auto()  # 4
+    HALF = enum.auto()  # 5
+    FLOAT = enum.auto()  # 6
+    DOUBLE = enum.auto()  # 7
+    COMPLEX32 = enum.auto()  # 8
+    COMPLEX64 = enum.auto()  # 9
+    COMPLEX128 = enum.auto()  # 10
+    BOOL = enum.auto()  # 11
+    QINT8 = enum.auto()  # 12
+    QUINT8 = enum.auto()  # 13
+    QINT32 = enum.auto()  # 14
+    BFLOAT16 = enum.auto()  # 15
+    FLOAT8E5M2 = enum.auto()  # 16
+    FLOAT8E4M3FN = enum.auto()  # 17
+    FLOAT8E5M2FNUZ = enum.auto()  # 18
+    FLOAT8E4M3FNUZ = enum.auto()  # 19
+    UNDEFINED = enum.auto()  # 20
+
+    @classmethod
+    def _from_name(cls, name: ScalarName | TorchName | str | None) -> JitScalarType:
+        """Convert a JIT scalar type or torch type name to ScalarType.
+
+        Note: DO NOT USE this API when `name` comes from a `torch._C.Value.type()` calls.
+            A "RuntimeError: INTERNAL ASSERT FAILED at "../aten/src/ATen/core/jit_type_base.h" can
+            be raised in several scenarios where shape info is not present.
+            Instead use `from_value` API which is safer.
+
+        Args:
+            name: JIT scalar type name (Byte) or torch type name (uint8_t).
+
+        Returns:
+            JitScalarType
+
+        Raises:
+           OnnxExporterError: if name is not a valid scalar type name or if it is None.
+        """
+        if name is None:
+            raise errors.OnnxExporterError("Scalar type name cannot be None")
+        if valid_scalar_name(name):
+            return _SCALAR_NAME_TO_TYPE[name]  # type: ignore[index]
+        if valid_torch_name(name):
+            return _TORCH_NAME_TO_SCALAR_TYPE[name]  # type: ignore[index]
+
+        raise errors.OnnxExporterError(f"Unknown torch or scalar type: '{name}'")
+
+    @classmethod
+    def from_dtype(cls, dtype: torch.dtype | None) -> JitScalarType:
+        """Convert a torch dtype to JitScalarType.
+
+        Note: DO NOT USE this API when `dtype` comes from a `torch._C.Value.type()` calls.
+            A "RuntimeError: INTERNAL ASSERT FAILED at "../aten/src/ATen/core/jit_type_base.h" can
+            be raised in several scenarios where shape info is not present.
+            Instead use `from_value` API which is safer.
+
+        Args:
+            dtype: A torch.dtype to create a JitScalarType from
+
+        Returns:
+            JitScalarType
+
+        Raises:
+            OnnxExporterError: if dtype is not a valid torch.dtype or if it is None.
+        """
+        if dtype not in _DTYPE_TO_SCALAR_TYPE:
+            raise errors.OnnxExporterError(f"Unknown dtype: {dtype}")
+        # pyrefly: ignore [bad-index]
+        return _DTYPE_TO_SCALAR_TYPE[dtype]
+
+    @classmethod
+    def from_onnx_type(
+        cls, onnx_type: int | _C_onnx.TensorProtoDataType | None
+    ) -> JitScalarType:
+        """Convert a ONNX data type to JitScalarType.
+
+        Args:
+            onnx_type: A torch._C._onnx.TensorProtoDataType to create a JitScalarType from
+
+        Returns:
+            JitScalarType
+
+        Raises:
+            OnnxExporterError: if dtype is not a valid torch.dtype or if it is None.
+        """
+        if onnx_type not in _ONNX_TO_SCALAR_TYPE:
+            raise errors.OnnxExporterError(f"Unknown onnx_type: {onnx_type}")
+        # pyrefly: ignore [redundant-cast]
+        return _ONNX_TO_SCALAR_TYPE[typing.cast(_C_onnx.TensorProtoDataType, onnx_type)]
+
+    @classmethod
+    def from_value(
+        cls, value: torch._C.Value | torch.Tensor | None, default=None
+    ) -> JitScalarType:
+        """Create a JitScalarType from an value's scalar type.
+
+        Args:
+            value: An object to fetch scalar type from.
+            default: The JitScalarType to return if a valid scalar cannot be fetched from value
+
+        Returns:
+            JitScalarType.
+
+        Raises:
+            OnnxExporterError: if value does not have a valid scalar type and default is None.
+            SymbolicValueError: when value.type()'s info are empty and default is None
+        """
+
+        if not isinstance(value, (torch._C.Value, torch.Tensor)) or (
+            isinstance(value, torch._C.Value) and value.node().mustBeNone()
+        ):
+            # default value of type JitScalarType is returned when value is not valid
+            if default is None:
+                raise errors.OnnxExporterError(
+                    "value must be either torch._C.Value or torch.Tensor objects."
+                )
+            elif not isinstance(default, JitScalarType):
+                raise errors.OnnxExporterError(
+                    "default value must be a JitScalarType object."
+                )
+            return default
+
+        # Each value type has their own way of storing scalar type
+        if isinstance(value, torch.Tensor):
+            return cls.from_dtype(value.dtype)
+        if isinstance(value.type(), torch.ListType):
+            try:
+                return cls.from_dtype(value.type().getElementType().dtype())
+            except RuntimeError:
+                return cls._from_name(str(value.type().getElementType()))
+        if isinstance(value.type(), torch._C.OptionalType):
+            if value.type().getElementType().dtype() is None:
+                if isinstance(default, JitScalarType):
+                    return default
+                raise errors.OnnxExporterError(
+                    "default value must be a JitScalarType object."
+                )
+            return cls.from_dtype(value.type().getElementType().dtype())
+
+        scalar_type = None
+        if value.node().kind() != "prim::Constant" or not isinstance(
+            value.type(), torch._C.NoneType
+        ):
+            # value must be a non-list torch._C.Value scalar
+            scalar_type = value.type().scalarType()
+
+        if scalar_type is not None:
+            return cls._from_name(scalar_type)
+
+        # When everything fails... try to default
+        if default is not None:
+            return default
+        raise errors.SymbolicValueError(
+            f"Cannot determine scalar type for this '{type(value.type())}' instance and "
+            "a default value was not provided.",
+            value,
+        )
+
+    def scalar_name(self) -> ScalarName:
+        """Convert a JitScalarType to a JIT scalar type name."""
+        return _SCALAR_TYPE_TO_NAME[self]
+
+    def torch_name(self) -> TorchName:
+        """Convert a JitScalarType to a torch type name."""
+        return _SCALAR_TYPE_TO_TORCH_NAME[self]
+
+    def dtype(self) -> torch.dtype:
+        """Convert a JitScalarType to a torch dtype."""
+        return _SCALAR_TYPE_TO_DTYPE[self]
+
+    def onnx_type(self) -> _C_onnx.TensorProtoDataType:
+        """Convert a JitScalarType to an ONNX data type."""
+        if self not in _SCALAR_TYPE_TO_ONNX:
+            raise errors.OnnxExporterError(
+                f"Scalar type {self} cannot be converted to ONNX"
+            )
+        return _SCALAR_TYPE_TO_ONNX[self]
+
+    def onnx_compatible(self) -> bool:
+        """Return whether this JitScalarType is compatible with ONNX."""
+        return (
+            self in _SCALAR_TYPE_TO_ONNX
+            and self != JitScalarType.UNDEFINED
+            and self != JitScalarType.COMPLEX32
+        )
+
+
+def valid_scalar_name(scalar_name: ScalarName | str) -> bool:
+    """Return whether the given scalar name is a valid JIT scalar type name."""
+    return scalar_name in _SCALAR_NAME_TO_TYPE
+
+
+def valid_torch_name(torch_name: TorchName | str) -> bool:
+    """Return whether the given torch name is a valid torch type name."""
+    return torch_name in _TORCH_NAME_TO_SCALAR_TYPE
+
+
+# https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h
+_SCALAR_TYPE_TO_NAME: dict[JitScalarType, ScalarName] = {
+    JitScalarType.BOOL: "Bool",
+    JitScalarType.UINT8: "Byte",
+    JitScalarType.INT8: "Char",
+    JitScalarType.INT16: "Short",
+    JitScalarType.INT: "Int",
+    JitScalarType.INT64: "Long",
+    JitScalarType.HALF: "Half",
+    JitScalarType.FLOAT: "Float",
+    JitScalarType.DOUBLE: "Double",
+    JitScalarType.COMPLEX32: "ComplexHalf",
+    JitScalarType.COMPLEX64: "ComplexFloat",
+    JitScalarType.COMPLEX128: "ComplexDouble",
+    JitScalarType.QINT8: "QInt8",
+    JitScalarType.QUINT8: "QUInt8",
+    JitScalarType.QINT32: "QInt32",
+    JitScalarType.BFLOAT16: "BFloat16",
+    JitScalarType.FLOAT8E5M2: "Float8E5M2",
+    JitScalarType.FLOAT8E4M3FN: "Float8E4M3FN",
+    JitScalarType.FLOAT8E5M2FNUZ: "Float8E5M2FNUZ",
+    JitScalarType.FLOAT8E4M3FNUZ: "Float8E4M3FNUZ",
+    JitScalarType.UNDEFINED: "Undefined",
+}
+
+_SCALAR_NAME_TO_TYPE: dict[ScalarName, JitScalarType] = {
+    v: k for k, v in _SCALAR_TYPE_TO_NAME.items()
+}
+
+_SCALAR_TYPE_TO_TORCH_NAME: dict[JitScalarType, TorchName] = {
+    JitScalarType.BOOL: "bool",
+    JitScalarType.UINT8: "uint8_t",
+    JitScalarType.INT8: "int8_t",
+    JitScalarType.INT16: "int16_t",
+    JitScalarType.INT: "int",
+    JitScalarType.INT64: "int64_t",
+    JitScalarType.HALF: "half",
+    JitScalarType.FLOAT: "float",
+    JitScalarType.DOUBLE: "double",
+    JitScalarType.COMPLEX32: "complex32",
+    JitScalarType.COMPLEX64: "complex64",
+    JitScalarType.COMPLEX128: "complex128",
+    JitScalarType.QINT8: "qint8",
+    JitScalarType.QUINT8: "quint8",
+    JitScalarType.QINT32: "qint32",
+    JitScalarType.BFLOAT16: "bfloat16",
+    JitScalarType.FLOAT8E5M2: "float8_e5m2",
+    JitScalarType.FLOAT8E4M3FN: "float8_e4m3fn",
+    JitScalarType.FLOAT8E5M2FNUZ: "float8_e5m2fnuz",
+    JitScalarType.FLOAT8E4M3FNUZ: "float8_e4m3fnuz",
+}
+
+_TORCH_NAME_TO_SCALAR_TYPE: dict[TorchName, JitScalarType] = {
+    v: k for k, v in _SCALAR_TYPE_TO_TORCH_NAME.items()
+}
+
+_SCALAR_TYPE_TO_ONNX = {
+    JitScalarType.BOOL: _C_onnx.TensorProtoDataType.BOOL,
+    JitScalarType.UINT8: _C_onnx.TensorProtoDataType.UINT8,
+    JitScalarType.INT8: _C_onnx.TensorProtoDataType.INT8,
+    JitScalarType.INT16: _C_onnx.TensorProtoDataType.INT16,
+    JitScalarType.INT: _C_onnx.TensorProtoDataType.INT32,
+    JitScalarType.INT64: _C_onnx.TensorProtoDataType.INT64,
+    JitScalarType.HALF: _C_onnx.TensorProtoDataType.FLOAT16,
+    JitScalarType.FLOAT: _C_onnx.TensorProtoDataType.FLOAT,
+    JitScalarType.DOUBLE: _C_onnx.TensorProtoDataType.DOUBLE,
+    JitScalarType.COMPLEX64: _C_onnx.TensorProtoDataType.COMPLEX64,
+    JitScalarType.COMPLEX128: _C_onnx.TensorProtoDataType.COMPLEX128,
+    JitScalarType.BFLOAT16: _C_onnx.TensorProtoDataType.BFLOAT16,
+    JitScalarType.UNDEFINED: _C_onnx.TensorProtoDataType.UNDEFINED,
+    JitScalarType.COMPLEX32: _C_onnx.TensorProtoDataType.UNDEFINED,
+    JitScalarType.QINT8: _C_onnx.TensorProtoDataType.INT8,
+    JitScalarType.QUINT8: _C_onnx.TensorProtoDataType.UINT8,
+    JitScalarType.QINT32: _C_onnx.TensorProtoDataType.INT32,
+    JitScalarType.FLOAT8E5M2: _C_onnx.TensorProtoDataType.FLOAT8E5M2,
+    JitScalarType.FLOAT8E4M3FN: _C_onnx.TensorProtoDataType.FLOAT8E4M3FN,
+    JitScalarType.FLOAT8E5M2FNUZ: _C_onnx.TensorProtoDataType.FLOAT8E5M2FNUZ,
+    JitScalarType.FLOAT8E4M3FNUZ: _C_onnx.TensorProtoDataType.FLOAT8E4M3FNUZ,
+}
+
+_ONNX_TO_SCALAR_TYPE = {v: k for k, v in _SCALAR_TYPE_TO_ONNX.items()}
+
+# source of truth is
+# https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_dtypes.cpp
+_SCALAR_TYPE_TO_DTYPE = {
+    JitScalarType.BOOL: torch.bool,
+    JitScalarType.UINT8: torch.uint8,
+    JitScalarType.INT8: torch.int8,
+    JitScalarType.INT16: torch.short,
+    JitScalarType.INT: torch.int,
+    JitScalarType.INT64: torch.int64,
+    JitScalarType.HALF: torch.half,
+    JitScalarType.FLOAT: torch.float,
+    JitScalarType.DOUBLE: torch.double,
+    JitScalarType.COMPLEX32: torch.complex32,
+    JitScalarType.COMPLEX64: torch.complex64,
+    JitScalarType.COMPLEX128: torch.complex128,
+    JitScalarType.QINT8: torch.qint8,
+    JitScalarType.QUINT8: torch.quint8,
+    JitScalarType.QINT32: torch.qint32,
+    JitScalarType.BFLOAT16: torch.bfloat16,
+    JitScalarType.FLOAT8E5M2: torch.float8_e5m2,
+    JitScalarType.FLOAT8E4M3FN: torch.float8_e4m3fn,
+    JitScalarType.FLOAT8E5M2FNUZ: torch.float8_e5m2fnuz,
+    JitScalarType.FLOAT8E4M3FNUZ: torch.float8_e4m3fnuz,
+}
+
+_DTYPE_TO_SCALAR_TYPE = {v: k for k, v in _SCALAR_TYPE_TO_DTYPE.items()}

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/jit_utils.py

@@ -0,0 +1,374 @@
+# mypy: allow-untyped-defs
+"""Utilities for manipulating the torch.Graph object and the torchscript."""
+
+from __future__ import annotations
+
+import dataclasses
+import re
+import typing
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import torch
+from torch import _C
+from torch.onnx._internal.torchscript_exporter import registration
+from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
+
+
+_ATTR_PATTERN = re.compile("^(.+)_(([ifstgz])|(ty))$")
+_SKIP_NODE_ATTRIBUTES = {"inplace", "aten"}
+
+
+@dataclasses.dataclass
+class GraphContext:
+    """Extra context for symbolic functions with all methods from torch.Graph.
+
+    NOTE: This class is not meant for external consumption. Please do not depend on
+    it outside of torch.onnx as the interface may evolve.
+
+    Attributes:
+        graph: The _C.Graph being constructed.
+        block: The current _C.Block being constructed.
+        opset: The opset version.
+        original_node: Current node that is being converted from.
+        params_dict: Mapping from graph initializer name to IValue.
+        env: Mapping from Torch domain graph Value to ONNX domain graph Value.
+        values_in_env: Set of all values in env, for constant-time lookups.
+        new_nodes: List that tracks all new nodes that are added (used to make
+            sure metadata is propagated to all new nodes).
+    """
+
+    graph: _C.Graph
+    block: _C.Block
+    opset: int
+    original_node: _C.Node
+    params_dict: dict[str, _C.IValue]
+    env: dict[_C.Value, _C.Value]
+    values_in_env: set[_C.Value]
+    new_nodes: list[_C.Node] = dataclasses.field(default_factory=list)
+
+    # Relay methods from _C.Graph for compatibility with symbolic functions that expect
+    # a _C.Graph
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.graph, name)
+
+    def op(
+        self,
+        opname: str,
+        *raw_args: torch.Tensor | _C.Value,
+        outputs: int = 1,
+        **kwargs,
+    ):
+        """Creates an ONNX operator "opname", taking "raw_args" as inputs and "kwargs" as attributes.
+
+        The set of operators and the inputs/attributes they take
+        is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md
+
+        Args:
+            opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified
+                with a namespace, e.g., `aten::add`.
+            raw_args: The inputs to the operator; usually provided
+                as arguments to the `symbolic` definition.
+            outputs: The number of outputs this operator returns.
+                By default an operator is assumed to return a single output.
+                If `outputs` is greater than one, this functions returns a tuple
+                of output `Value`, representing each output of the ONNX operator
+                in order.
+            kwargs: The attributes of the ONNX operator, whose keys are named
+                according to the following convention: `alpha_f` indicates
+                the `alpha` attribute with type `f`.  The valid type specifiers are
+                `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
+                specified with type float accepts either a single float, or a
+                list of floats (e.g., you would say `dims_i` for a `dims` attribute
+                that takes a list of integers).
+
+        Returns:
+            The value representing the single output of this operator (see the `outputs`
+            keyword argument for multi-return nodes).
+        """
+        return _add_op(self, opname, *raw_args, outputs=outputs, **kwargs)
+
+    def aten_op(self, operator: str, *args, overload_name: str = "", **kwargs):
+        """Generates an ONNX ATen op node.
+
+        This function is for backward compatibility with the old symbolic functions.
+        """
+        return self.op(
+            "aten::ATen",
+            *args,
+            operator_s=operator,
+            overload_name_s=overload_name,
+            **kwargs,
+        )
+
+    # NOTE: For backward compatibility with the old symbolic functions.
+    # We are probably going to remove this only after the fx exporter is established.
+    at = aten_op
+
+    def onnxscript_op(
+        self,
+        onnx_fn,
+        *raw_args: torch.Tensor | _C.Value,
+        outputs: int = 1,
+        **kwargs,
+    ):
+        """Creates an ONNX operator from onnx-script function, taking "raw_args" as inputs and "kwargs" as attributes.
+
+        onnx-script repository: https://github.com/microsoft/onnx-script
+
+        Args:
+            onnx_fn: ONNXFunction from onnx-script; An example can be found at
+                https://github.com/microsoft/onnx-script#example
+            raw_args: The inputs to the operator; usually provided
+                as arguments to the `symbolic` definition.
+            outputs: The number of outputs this operator returns.
+                By default an operator is assumed to return a single output.
+                If `outputs` is greater than one, this functions returns a tuple
+                of output `Value`, representing each output of the ONNX operator
+                in order.
+            kwargs: The attributes of the ONNX operator, whose keys are named
+                according to the following convention: `alpha_f` indicates
+                the `alpha` attribute with type `f`.  The valid type specifiers are
+                `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
+                specified with type float accepts either a single float, or a
+                list of floats (e.g., you would say `dims_i` for a `dims` attribute
+                that takes a list of integers).
+
+        Returns:
+            The value representing the single output of this operator (see the `outputs`
+            keyword argument for multi-return nodes).
+        """
+        # NOTE(titaiwang): This is using class attributes, and it needs to be updated
+        # if onnx-script makes any change on these.
+        symbolic_name = f"{onnx_fn.opset.domain}::{onnx_fn.name}"
+        opset_version = onnx_fn.opset.version
+
+        registration.custom_onnx_symbolic(symbolic_name, opset_version)(onnx_fn)
+
+        return _add_op(self, symbolic_name, *raw_args, outputs=outputs, **kwargs)
+
+
+def add_op_with_blocks(
+    graph_context: GraphContext,
+    opname: str,
+    *inputs: _C.Value,
+    outputs: int = 1,
+    n_blocks: int = 1,
+    **attributes,
+) -> tuple[Any, tuple[GraphContext, ...], _C.Node]:
+    """Creates an ONNX operator "opname", taking inputs and attributes.
+
+    Args:
+        graph_context: The context for the current graph.
+        opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified
+            with a namespace, e.g., `aten::add`.
+        inputs: The inputs to the operator.
+        outputs: The number of outputs this operator returns.
+            By default an operator is assumed to return a single output.
+            If `outputs` is greater than one, this functions returns a tuple
+            of output `Value`, representing each output of the ONNX operator
+            in order.
+        n_blocks: The number of sub-blocks to create in the node.
+        attributes: The attributes of the ONNX operator.
+
+    Returns:
+        A tuple of (output_values, new_contexts, node) where:
+            output_values: One or more output value of this operator
+                (see the `outputs` keyword argument for multi-return nodes).
+            new_contexts: A tuple of new graph contexts for each sub-block.
+            node: The node representing the operator.
+    """
+
+    output_values = graph_context.op(opname, *inputs, outputs=outputs, **attributes)
+    if isinstance(output_values, Sequence):
+        node = output_values[0].node()
+    else:
+        node = output_values.node()
+
+    new_contexts = []
+    for _ in range(n_blocks):
+        new_block = node.addBlock()
+        # Create shallow copy of the graph context and update the block
+        new_context = dataclasses.replace(graph_context, block=new_block)
+        new_contexts.append(new_context)
+
+    return output_values, tuple(new_contexts), node
+
+
+def _add_op(
+    graph_context: GraphContext,
+    opname: str,
+    *args: torch.Tensor | _C.Value,
+    outputs: int = 1,
+    **kwargs,
+):
+    """Creates an ONNX operator "opname", taking "args" as inputs and attributes "kwargs".
+
+    The set of operators and the inputs/attributes they take
+    is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md
+
+    Args:
+        graph_context: The Torch Graph or Block.
+        opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified
+            with a namespace, e.g., `aten::add`.
+        args: The inputs to the operator; usually provided
+            as arguments to the `symbolic` definition.
+        outputs: The number of outputs this operator returns.
+            By default an operator is assumed to return a single output.
+            If `outputs` is greater than one, this functions returns a tuple
+            of output `Value`, representing each output of the ONNX operator
+            in order.
+        kwargs: The attributes of the ONNX operator, whose keys are named
+            according to the following convention: `alpha_f` indicates
+            the `alpha` attribute with type `f`.  The valid type specifiers are
+            `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
+            specified with type float accepts either a single float, or a
+            list of floats (e.g., you would say `dims_i` for a `dims` attribute
+            that takes a list of integers).
+
+    Returns:
+        (Union[_C.Value, Tuple[_C.Value, ...]])
+        The value representing the single output of this operator (see the `outputs`
+        keyword argument for multi-return nodes).
+    """
+    inputs = [_const_if_tensor(graph_context, arg) for arg in args]
+    # Filter out None attributes, this can be convenient client side because
+    # now they can pass through None attributes, and have them not show up
+    attributes = {k: v for k, v in kwargs.items() if v is not None}
+
+    if "::" not in opname:
+        opname = "onnx::" + opname
+
+    node = _create_node(
+        graph_context.block,
+        opname,
+        inputs,
+        attributes,
+        params_dict=graph_context.params_dict,
+        opset_version=graph_context.opset,
+        n_outputs=outputs,
+        shape_inference=GLOBALS.onnx_shape_inference,
+    )
+    graph_context.new_nodes.append(node)
+
+    if outputs == 1:
+        return node.output()
+    return tuple(node.outputs())
+
+
+def _const_if_tensor(graph_context: GraphContext, arg):
+    if arg is None:
+        return arg
+    if isinstance(arg, _C.Value):
+        return arg
+
+    return _add_op(graph_context, "onnx::Constant", value_z=arg)
+
+
+def _create_node(
+    graph_or_block: _C.Graph | _C.Block,
+    domain_op: str,
+    inputs: Sequence,
+    attributes: dict,
+    params_dict: dict,
+    opset_version: int,
+    n_outputs: int,
+    shape_inference: bool = True,
+) -> _C.Node:
+    """Creates an node 'domain_op', taking inputs and attributes."""
+    if isinstance(graph_or_block, _C.Graph):
+        graph = graph_or_block
+        node = graph.create(domain_op, inputs, n_outputs)
+        node = graph.insertNode(node)
+    elif isinstance(graph_or_block, _C.Block):
+        block = graph_or_block
+        node = block.addNode(domain_op, inputs)
+
+        # Block does not have create defined, so we need to add outputs manually
+        if n_outputs > 1:
+            for _ in range(1, n_outputs):
+                node.addOutput()
+
+    node_outputs = tuple(node.outputs())  # type: ignore[possibly-undefined]
+    if len(node_outputs) != n_outputs:
+        raise AssertionError(
+            f"len(node_outputs)={len(node_outputs)} != n_outputs={n_outputs}"
+        )
+
+    aten = domain_op.startswith("aten::")
+
+    # Add all attributes
+    for key, value in sorted(attributes.items()):
+        if key in _SKIP_NODE_ATTRIBUTES:
+            continue
+        # pyrefly: ignore [unbound-name]
+        _add_attribute(node, key, value, aten=aten)
+    if shape_inference:
+        # pyrefly: ignore [unbound-name]
+        _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version)
+    # pyrefly: ignore [unbound-name]
+    return node
+
+
+def _is_onnx_list(value):
+    return isinstance(value, Iterable) and not isinstance(
+        value, (str, bytes, torch.Tensor)
+    )
+
+
+def _scalar(x: torch.Tensor):
+    """Convert a scalar tensor into a Python value."""
+    if x.numel() != 1:
+        raise AssertionError(f"Expected numel() == 1, got {x.numel()}")
+    return x[0]
+
+
+def _add_attribute(node: _C.Node, key: str, value: Any, aten: bool):
+    r"""Initializes the right attribute based on type of value."""
+    m = _ATTR_PATTERN.match(key)
+    if m is None:
+        raise ValueError(
+            f"Invalid attribute specifier '{key}' names "
+            "must be suffixed with type, e.g. 'dim_i' or 'dims_i'"
+        )
+    name, kind = m.group(1), m.group(2)
+    if _is_onnx_list(value):
+        kind += "s"
+
+    return getattr(node, f"{kind}_")(name, value)
+
+
+def _is_tensor(x: _C.Value) -> bool:
+    return x.type().isSubtypeOf(_C.TensorType.get())
+
+
+def get_device_from_value(value: _C.Value) -> torch.device | None:
+    if not _is_tensor(value):
+        return None
+    tensor_type = typing.cast(_C.TensorType, value.type())
+    return tensor_type.device()
+
+
+def parse_node_kind(kind: str) -> tuple[str, str]:
+    """Parse node kind into domain and Op name."""
+    if "::" not in kind:
+        raise ValueError(f"Node kind: {kind} is invalid. '::' is not in node kind.")
+    domain, opname = kind.split("::", 1)
+    if "::" in opname:
+        raise ValueError(f"Node kind: {kind} is invalid. '::' should only appear once.")
+    return domain, opname
+
+
+def is_aten(domain: str) -> bool:
+    """Check if the domain is official."""
+    return domain == "aten"
+
+
+def is_prim(domain: str) -> bool:
+    """Check if the domain is official."""
+    return domain == "prim"
+
+
+def is_onnx(domain: str) -> bool:
+    """Check if the domain is official."""
+    return domain == "onnx"

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/onnx_proto_utils.py

@@ -0,0 +1,251 @@
+# mypy: allow-untyped-defs
+"""Utilities for manipulating the onnx and onnx-script dependencies and ONNX proto."""
+
+from __future__ import annotations
+
+import glob
+import os
+import shutil
+from typing import Any, TYPE_CHECKING
+
+import torch
+import torch.serialization
+from torch.onnx import errors
+from torch.onnx._internal.torchscript_exporter import jit_utils, registration
+
+
+if TYPE_CHECKING:
+    import io
+    from collections.abc import Mapping
+
+
+def export_as_test_case(
+    model_bytes: bytes, inputs_data, outputs_data, name: str, dir: str
+) -> str:
+    """Export an ONNX model as a self contained ONNX test case.
+
+    The test case contains the model and the inputs/outputs data. The directory structure
+    is as follows:
+
+    dir
+    \u251c\u2500\u2500 test_<name>
+    \u2502   \u251c\u2500\u2500 model.onnx
+    \u2502   \u2514\u2500\u2500 test_data_set_0
+    \u2502       \u251c\u2500\u2500 input_0.pb
+    \u2502       \u251c\u2500\u2500 input_1.pb
+    \u2502       \u251c\u2500\u2500 output_0.pb
+    \u2502       \u2514\u2500\u2500 output_1.pb
+
+    Args:
+        model_bytes: The ONNX model in bytes.
+        inputs_data: The inputs data, nested data structure of numpy.ndarray.
+        outputs_data: The outputs data, nested data structure of numpy.ndarray.
+
+    Returns:
+        The path to the test case directory.
+    """
+    try:
+        import onnx
+    except ImportError as exc:
+        raise ImportError(
+            "Export test case to ONNX format failed: Please install ONNX."
+        ) from exc
+
+    test_case_dir = os.path.join(dir, "test_" + name)
+    os.makedirs(test_case_dir, exist_ok=True)
+    _export_file(
+        model_bytes,
+        os.path.join(test_case_dir, "model.onnx"),
+        {},
+    )
+    data_set_dir = os.path.join(test_case_dir, "test_data_set_0")
+    if os.path.exists(data_set_dir):
+        shutil.rmtree(data_set_dir)
+    os.makedirs(data_set_dir)
+
+    proto = onnx.load_model_from_string(model_bytes)  # type: ignore[attr-defined]
+
+    for i, (input_proto, input) in enumerate(zip(proto.graph.input, inputs_data)):
+        export_data(input, input_proto, os.path.join(data_set_dir, f"input_{i}.pb"))
+    for i, (output_proto, output) in enumerate(zip(proto.graph.output, outputs_data)):
+        export_data(output, output_proto, os.path.join(data_set_dir, f"output_{i}.pb"))
+
+    return test_case_dir
+
+
+def load_test_case(dir: str) -> tuple[bytes, Any, Any]:
+    """Load a self contained ONNX test case from a directory.
+
+    The test case must contain the model and the inputs/outputs data. The directory structure
+    should be as follows:
+
+    dir
+    \u251c\u2500\u2500 test_<name>
+    \u2502   \u251c\u2500\u2500 model.onnx
+    \u2502   \u2514\u2500\u2500 test_data_set_0
+    \u2502       \u251c\u2500\u2500 input_0.pb
+    \u2502       \u251c\u2500\u2500 input_1.pb
+    \u2502       \u251c\u2500\u2500 output_0.pb
+    \u2502       \u2514\u2500\u2500 output_1.pb
+
+    Args:
+        dir: The directory containing the test case.
+
+    Returns:
+        model_bytes: The ONNX model in bytes.
+        inputs: the inputs data, mapping from input name to numpy.ndarray.
+        outputs: the outputs data, mapping from output name to numpy.ndarray.
+    """
+    try:
+        import onnx
+        from onnx import numpy_helper  # type: ignore[attr-defined]
+    except ImportError as exc:
+        raise ImportError(
+            "Load test case from ONNX format failed: Please install ONNX."
+        ) from exc
+
+    with open(os.path.join(dir, "model.onnx"), "rb") as f:
+        model_bytes = f.read()
+
+    test_data_dir = os.path.join(dir, "test_data_set_0")
+
+    inputs = {}
+    input_files = glob.glob(os.path.join(test_data_dir, "input_*.pb"))
+    for input_file in input_files:
+        tensor = onnx.load_tensor(input_file)  # type: ignore[attr-defined]
+        inputs[tensor.name] = numpy_helper.to_array(tensor)
+    outputs = {}
+    output_files = glob.glob(os.path.join(test_data_dir, "output_*.pb"))
+    for output_file in output_files:
+        tensor = onnx.load_tensor(output_file)  # type: ignore[attr-defined]
+        outputs[tensor.name] = numpy_helper.to_array(tensor)
+
+    return model_bytes, inputs, outputs
+
+
+def export_data(data, value_info_proto, f: str) -> None:
+    """Export data to ONNX protobuf format.
+
+    Args:
+        data: The data to export, nested data structure of numpy.ndarray.
+        value_info_proto: The ValueInfoProto of the data. The type of the ValueInfoProto
+            determines how the data is stored.
+        f: The file to write the data to.
+    """
+    try:
+        from onnx import numpy_helper  # type: ignore[attr-defined]
+    except ImportError as exc:
+        raise ImportError(
+            "Export data to ONNX format failed: Please install ONNX."
+        ) from exc
+
+    with open(f, "wb") as opened_file:
+        if value_info_proto.type.HasField("map_type"):
+            opened_file.write(
+                numpy_helper.from_dict(data, value_info_proto.name).SerializeToString()
+            )
+        elif value_info_proto.type.HasField("sequence_type"):
+            opened_file.write(
+                numpy_helper.from_list(data, value_info_proto.name).SerializeToString()
+            )
+        elif value_info_proto.type.HasField("optional_type"):
+            opened_file.write(
+                numpy_helper.from_optional(
+                    data, value_info_proto.name
+                ).SerializeToString()
+            )
+        else:
+            if not value_info_proto.type.HasField("tensor_type"):
+                raise AssertionError("Expected tensor_type field to be set")
+            opened_file.write(
+                numpy_helper.from_array(data, value_info_proto.name).SerializeToString()
+            )
+
+
+def _export_file(
+    model_bytes: bytes,
+    f: io.BytesIO | str,
+    export_map: Mapping[str, bytes],
+) -> None:
+    """export/write model bytes into directory/protobuf/zip"""
+    if len(export_map) != 0:
+        raise AssertionError(f"export_map must be empty, got {len(export_map)} items")
+    with torch.serialization._open_file_like(f, "wb") as opened_file:
+        opened_file.write(model_bytes)
+
+
+def _add_onnxscript_fn(
+    model_bytes: bytes,
+    custom_opsets: Mapping[str, int],
+) -> bytes:
+    """Insert model-included custom onnx-script function into ModelProto"""
+    try:
+        import onnx
+    except ImportError as e:
+        raise errors.OnnxExporterError("Module onnx is not installed!") from e
+
+    # For > 2GB model, onnx.load_fromstring would fail. However, because
+    # in _export_onnx, the tensors should be saved separately if the proto
+    # size > 2GB, and if it for some reason did not, the model would fail on
+    # serialization anyway in terms of the protobuf limitation. So we don't
+    # need to worry about > 2GB model getting here.
+    model_proto = onnx.load_model_from_string(model_bytes)  # type: ignore[attr-defined]
+
+    # Iterate graph nodes to insert only the included custom
+    # function_proto into model_proto
+    onnx_function_list = []  # type: ignore[var-annotated]
+    included_node_func: set[str] = set()
+    # onnx_function_list and included_node_func are expanded in-place
+    _find_onnxscript_op(
+        model_proto.graph, included_node_func, custom_opsets, onnx_function_list
+    )
+
+    if onnx_function_list:
+        model_proto.functions.extend(onnx_function_list)
+        model_bytes = model_proto.SerializeToString()
+    return model_bytes
+
+
+def _find_onnxscript_op(
+    graph_proto,
+    included_node_func: set[str],
+    custom_opsets: Mapping[str, int],
+    onnx_function_list: list,
+):
+    """Recursively iterate ModelProto to find ONNXFunction op as it may contain control flow Op."""
+    for node in graph_proto.node:
+        node_kind = node.domain + "::" + node.op_type
+        # Recursive needed for control flow nodes: IF/Loop which has inner graph_proto
+        for attr in node.attribute:
+            if attr.g is not None:
+                _find_onnxscript_op(
+                    attr.g, included_node_func, custom_opsets, onnx_function_list
+                )
+        # Only custom Op with ONNX function and aten with symbolic_fn should be found in registry
+        onnx_function_group = registration.registry.get_function_group(node_kind)
+        # Ruled out corner cases: onnx/prim in registry
+        if (
+            node.domain
+            and not jit_utils.is_aten(node.domain)
+            and not jit_utils.is_prim(node.domain)
+            and not jit_utils.is_onnx(node.domain)
+            and onnx_function_group is not None
+            and node_kind not in included_node_func
+        ):
+            specified_version = custom_opsets.get(node.domain, 1)
+            onnx_fn = onnx_function_group.get(specified_version)
+            if onnx_fn is not None:
+                if hasattr(onnx_fn, "to_function_proto"):
+                    onnx_function_proto = onnx_fn.to_function_proto()  # type: ignore[attr-defined]
+                    onnx_function_list.append(onnx_function_proto)
+                    included_node_func.add(node_kind)
+                continue
+
+            raise errors.UnsupportedOperatorError(
+                node_kind,
+                specified_version,
+                onnx_function_group.get_min_supported()
+                if onnx_function_group
+                else None,
+            )
+    return onnx_function_list, included_node_func

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/registration.py

@@ -0,0 +1,337 @@
+# mypy: allow-untyped-defs
+"""Module for handling symbolic function registration."""
+
+import warnings
+from collections.abc import Callable, Collection, Sequence
+from typing import Generic, TypeVar
+from typing_extensions import ParamSpec
+
+from torch.onnx import _constants, errors
+
+
+OpsetVersion = int
+
+
+def _dispatch_opset_version(
+    target: OpsetVersion, registered_opsets: Collection[OpsetVersion]
+) -> OpsetVersion | None:
+    """Finds the registered opset given a target opset version and the available opsets.
+
+    Args:
+        target: The target opset version.
+        registered_opsets: The available opsets.
+
+    Returns:
+        The registered opset version.
+    """
+    if not registered_opsets:
+        return None
+
+    descending_registered_versions = sorted(registered_opsets, reverse=True)
+    # Linear search for the opset version, which is fine since the number of opset
+    # versions is small.
+
+    if target >= _constants.ONNX_BASE_OPSET:
+        # Always look down toward opset 1 when the target is >= ONNX_BASE_OPSET (opset 9).
+        # When a custom op is register at opset 1, we want to be able to discover it as a
+        # fallback for all opsets >= ONNX_BASE_OPSET.
+        for version in descending_registered_versions:
+            if version <= target:
+                return version
+        return None
+
+    # target < opset 9. This is the legacy behavior to support opset 7 and opset 8.
+    # for caffe2 support. We search up toward opset 9.
+    for version in reversed(descending_registered_versions):
+        # Count back up until _constants.ONNX_BASE_OPSET
+        if target <= version <= _constants.ONNX_BASE_OPSET:
+            return version
+
+    return None
+
+
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
+
+
+class OverrideDict(Collection[_K], Generic[_K, _V]):
+    """A dictionary that merges built-in and custom symbolic functions.
+
+    It supports overriding and un-overriding built-in symbolic functions with custom
+    ones.
+    """
+
+    def __init__(self) -> None:
+        self._base: dict[_K, _V] = {}
+        self._overrides: dict[_K, _V] = {}
+        self._merged: dict[_K, _V] = {}
+
+    def set_base(self, key: _K, value: _V) -> None:
+        self._base[key] = value
+        if key not in self._overrides:
+            self._merged[key] = value
+
+    def in_base(self, key: _K) -> bool:
+        """Checks if a key is in the base dictionary."""
+        return key in self._base
+
+    def override(self, key: _K, value: _V) -> None:
+        """Overrides a base key-value with a new pair."""
+        self._overrides[key] = value
+        self._merged[key] = value
+
+    def remove_override(self, key: _K) -> None:
+        """Un-overrides a key-value pair."""
+        self._overrides.pop(key, None)  # type: ignore[arg-type]
+        self._merged.pop(key, None)  # type: ignore[arg-type]
+        if key in self._base:
+            self._merged[key] = self._base[key]
+
+    def overridden(self, key: _K) -> bool:
+        """Checks if a key-value pair is overridden."""
+        return key in self._overrides
+
+    def __getitem__(self, key: _K) -> _V:
+        return self._merged[key]
+
+    def get(self, key: _K, default: _V | None = None):
+        return self._merged.get(key, default)
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._merged
+
+    def __iter__(self):
+        return iter(self._merged)
+
+    def __len__(self) -> int:
+        return len(self._merged)
+
+    def __repr__(self) -> str:
+        return f"OverrideDict(base={self._base}, overrides={self._overrides})"
+
+    def __bool__(self) -> bool:
+        return bool(self._merged)
+
+
+class _SymbolicFunctionGroup:
+    """Different versions of symbolic functions registered to the same name.
+
+    O(number of registered versions of an op) search is performed to find the most
+    recent version of the op.
+
+    The registration is delayed until op is used to improve startup time.
+
+    Function overloads with different arguments are not allowed.
+    Custom op overrides are supported.
+    """
+
+    def __init__(self, name: str) -> None:
+        self._name = name
+        # A dictionary of functions, keyed by the opset version.
+        self._functions: OverrideDict[OpsetVersion, Callable] = OverrideDict()
+
+    def __repr__(self) -> str:
+        return f"_SymbolicFunctionGroup({self._name}, registered={self._functions})"
+
+    def __getitem__(self, key: OpsetVersion) -> Callable:
+        result = self.get(key)
+        if result is None:
+            raise KeyError(key)
+        return result
+
+    # TODO(justinchuby): Add @functools.lru_cache(maxsize=None) if lookup time becomes
+    # a problem.
+    def get(self, opset: OpsetVersion) -> Callable | None:
+        """Find the most recent version of the function."""
+        version = _dispatch_opset_version(opset, self._functions)
+        if version is None:
+            return None
+
+        return self._functions[version]
+
+    def add(self, func: Callable, opset: OpsetVersion) -> None:
+        """Adds a symbolic function.
+
+        Args:
+            func: The function to add.
+            opset: The opset version of the function to add.
+        """
+        if self._functions.in_base(opset):
+            warnings.warn(
+                f"Symbolic function '{self._name}' already registered for opset {opset}. "
+                f"Replacing the existing function with new function. This is unexpected. "
+                f"Please report it on {_constants.PYTORCH_GITHUB_ISSUES_URL}.",
+                errors.OnnxExporterWarning,
+                stacklevel=2,
+            )
+        self._functions.set_base(opset, func)
+
+    def add_custom(self, func: Callable, opset: OpsetVersion) -> None:
+        """Adds a custom symbolic function.
+
+        Args:
+            func: The symbolic function to register.
+            opset: The corresponding opset version.
+        """
+        self._functions.override(opset, func)
+
+    def remove_custom(self, opset: OpsetVersion) -> None:
+        """Removes a custom symbolic function.
+
+        Args:
+            opset: The opset version of the custom function to remove.
+        """
+        if not self._functions.overridden(opset):
+            warnings.warn(
+                f"No custom function registered for '{self._name}' opset {opset}",
+                stacklevel=2,
+            )
+            return
+        self._functions.remove_override(opset)
+
+    def get_min_supported(self) -> OpsetVersion:
+        """Returns the lowest built-in opset version supported by the function."""
+        return min(self._functions)
+
+
+class SymbolicRegistry:
+    """Registry for symbolic functions.
+
+    The registry maintains a mapping from qualified names to symbolic functions.
+    It is used to register new symbolic functions and to dispatch calls to
+    the appropriate function.
+    """
+
+    def __init__(self) -> None:
+        self._registry: dict[str, _SymbolicFunctionGroup] = {}
+
+    def register(
+        self, name: str, opset: OpsetVersion, func: Callable, custom: bool = False
+    ) -> None:
+        """Registers a symbolic function.
+
+        Args:
+            name: The qualified name of the function to register. In the form of 'domain::op'.
+                E.g. 'aten::add'.
+            opset: The opset version of the function to register.
+            func: The symbolic function to register.
+            custom: Whether the function is a custom function that overrides existing ones.
+
+        Raises:
+            ValueError: If the separator '::' is not in the name.
+        """
+        if "::" not in name:
+            raise ValueError(
+                f"The name must be in the form of 'domain::op', not '{name}'"
+            )
+        symbolic_functions = self._registry.setdefault(
+            name, _SymbolicFunctionGroup(name)
+        )
+        if custom:
+            symbolic_functions.add_custom(func, opset)
+        else:
+            symbolic_functions.add(func, opset)
+
+    def unregister(self, name: str, opset: OpsetVersion) -> None:
+        """Unregisters a symbolic function.
+
+        Args:
+            name: The qualified name of the function to unregister.
+            opset: The opset version of the function to unregister.
+        """
+        if name not in self._registry:
+            return
+        self._registry[name].remove_custom(opset)
+
+    def get_function_group(self, name: str) -> _SymbolicFunctionGroup | None:
+        """Returns the function group for the given name."""
+        return self._registry.get(name)
+
+    def is_registered_op(self, name: str, version: int) -> bool:
+        """Returns whether the given op is registered for the given opset version."""
+        functions = self.get_function_group(name)
+        if functions is None:
+            return False
+        return functions.get(version) is not None
+
+    def all_functions(self) -> set[str]:
+        """Returns the set of all registered function names."""
+        return set(self._registry)
+
+
+def onnx_symbolic(
+    name: str,
+    opset: OpsetVersion | Sequence[OpsetVersion],
+    decorate: Sequence[Callable] | None = None,
+    custom: bool = False,
+) -> Callable:
+    """Registers a symbolic function.
+
+    Usage::
+
+    ```
+    @onnx_symbolic(
+        "aten::symbolic_b",
+        opset=10,
+        decorate=[quantized_aten_handler(scale=1 / 128, zero_point=0)],
+    )
+    @symbolic_helper.parse_args("v", "v", "b")
+    def symbolic_b(g: _C.Graph, x: _C.Value, y: _C.Value, arg1: bool) -> _C.Value: ...
+    ```
+
+    Args:
+        name: The qualified name of the function in the form of 'domain::op'.
+            E.g. 'aten::add'.
+        opset: The opset versions of the function to register at.
+        decorate: A sequence of decorators to apply to the function.
+        custom: Whether the function is a custom symbolic function.
+
+    Raises:
+        ValueError: If the separator '::' is not in the name.
+    """
+
+    def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]:
+        decorated = func
+        if decorate is not None:
+            for decorate_func in decorate:
+                decorated = decorate_func(decorated)
+
+        global registry
+        nonlocal opset
+        if isinstance(opset, OpsetVersion):
+            opset = (opset,)
+        for opset_version in opset:
+            registry.register(name, opset_version, decorated, custom=custom)
+
+        # Return the original function because the decorators in "decorate" are only
+        # specific to the instance being registered.
+        return func
+
+    return wrapper
+
+
+def custom_onnx_symbolic(
+    name: str,
+    opset: OpsetVersion | Sequence[OpsetVersion],
+    decorate: Sequence[Callable] | None = None,
+) -> Callable:
+    """Registers a custom symbolic function.
+
+    Args:
+        name: the qualified name of the function.
+        opset: the opset version of the function.
+        decorate: a sequence of decorators to apply to the function.
+
+    Returns:
+        The decorator.
+
+    Raises:
+        ValueError: If the separator '::' is not in the name.
+    """
+    return onnx_symbolic(name, opset, decorate, custom=True)
+
+
+# The registry for all symbolic functions.
+registry = SymbolicRegistry()

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset12.py

@@ -0,0 +1,469 @@
+# mypy: allow-untyped-defs
+# mypy: disable-error-code=arg-type
+from __future__ import annotations
+
+import functools
+import sys
+
+import torch
+from torch._C import _onnx as _C_onnx
+from torch.onnx import errors
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+    symbolic_opset9 as opset9,
+    utils,
+)
+
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in README.md
+
+# This file exports ONNX ops for opset 12
+
+__all__ = [
+    "argmax",
+    "argmin",
+    "binary_cross_entropy_with_logits",
+    "celu",
+    "cross_entropy_loss",
+    "dropout",
+    "einsum",
+    "ge",
+    "le",
+    "native_dropout",
+    "nll_loss",
+    "nll_loss2d",
+    "nll_loss_nd",
+    "outer",
+    "pow",
+    "tensordot",
+    "unfold",
+]
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=12)
+
+
+def _einsum_helper(g: jit_utils.GraphContext, equation, tensors):
+    if not tensors:
+        raise RuntimeError("Einsum inputs are empty.")
+    # ONNX does not support bool for Einsum inputs.
+    if symbolic_helper._is_bool(tensors[0]):
+        tensors = [
+            g.op("Cast", tensor, to_i=_C_onnx.TensorProtoDataType.INT64)
+            for tensor in tensors
+        ]
+        return g.op(
+            "Cast",
+            g.op("Einsum", *tensors, equation_s=equation),
+            to_i=_C_onnx.TensorProtoDataType.BOOL,
+        )
+    else:
+        return g.op("Einsum", *tensors, equation_s=equation)
+
+
+@_onnx_symbolic("aten::einsum")
+@symbolic_helper.parse_args("s", "v", "is")
+def einsum(g: jit_utils.GraphContext, equation, tensor_list, path=None):
+    tensors = symbolic_helper._unpack_list(tensor_list)
+    return _einsum_helper(g, equation, tensors)
+
+
+@_onnx_symbolic("aten::outer")
+@symbolic_helper.parse_args("v", "v")
+def outer(g: jit_utils.GraphContext, input, other):
+    # make sure to cast other to self's type
+    if _type_utils.JitScalarType.from_value(
+        other, _type_utils.JitScalarType.UNDEFINED
+    ) != _type_utils.JitScalarType.from_value(input):
+        other = g.op(
+            "Cast",
+            other,
+            to_i=_type_utils.JitScalarType.from_value(input).onnx_type(),
+        )
+    return _einsum_helper(g, "i,j->ij", [input, other])
+
+
+def _dropout_returns_masked_input_and_mask(
+    g: jit_utils.GraphContext, input: torch._C.Value, p: float, train: bool
+) -> tuple[torch._C.Value, torch._C.Value | None]:
+    symbolic_helper.check_training_mode(train, "dropout")
+    # In eval mode, dropout is non-op. That is, if the node's
+    # train param is set to False, dropout just returns its inputs.
+    if not train:
+        return input, None
+    p = g.op("Constant", value_t=torch.tensor(p))
+    t = g.op("Constant", value_t=torch.tensor(train, dtype=torch.bool))
+    r, mask = g.op("Dropout", input, p, t, outputs=2)
+    return r, mask
+
+
+@_onnx_symbolic("aten::dropout")
+@symbolic_helper.parse_args("v", "f", "b")
+def dropout(g: jit_utils.GraphContext, input, p, train):
+    masked, _ = _dropout_returns_masked_input_and_mask(g, input, p, train)
+    return masked
+
+
+@_onnx_symbolic("aten::native_dropout")
+@symbolic_helper.parse_args("v", "f", "b")
+def native_dropout(g: jit_utils.GraphContext, input, p, train):
+    return _dropout_returns_masked_input_and_mask(g, input, p, train)
+
+
+@_onnx_symbolic("aten::nll_loss")
+def nll_loss(g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index):
+    # none reduction : onnx::Constant[value={0}]
+    # mean reduction : onnx::Constant[value={1}]
+    # sum reduction : onnx::Constant[value={2}]
+    reduction = symbolic_helper._maybe_get_const(reduction, "i")
+    reduction_vals = ["none", "mean", "sum"]
+    reduction = reduction_vals[reduction]
+
+    # in onnx NegativeLogLikelihoodLoss specification, ignore_index is optional without default value.
+    # therefore we need to set ignore_index attribute even if it is not specified (e.g. ignore_index=-100).
+    ignore_index = symbolic_helper._maybe_get_const(ignore_index, "i")
+    if weight.node().mustBeNone():
+        nllloss = g.op(
+            "NegativeLogLikelihoodLoss",
+            self,
+            target,
+            reduction_s=reduction,
+            ignore_index_i=ignore_index,
+        )
+    else:
+        nllloss = g.op(
+            "NegativeLogLikelihoodLoss",
+            self,
+            target,
+            weight,
+            reduction_s=reduction,
+            ignore_index_i=ignore_index,
+        )
+
+    return nllloss
+
+
+@_onnx_symbolic("aten::nll_loss2d")
+def nll_loss2d(
+    g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index
+):
+    return nll_loss(g, self, target, weight, reduction, ignore_index)
+
+
+@_onnx_symbolic("aten::nll_loss_nd")
+def nll_loss_nd(
+    g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index
+):
+    return nll_loss(g, self, target, weight, reduction, ignore_index)
+
+
+@_onnx_symbolic("aten::cross_entropy_loss")
+def cross_entropy_loss(
+    g: jit_utils.GraphContext,
+    self,
+    target,
+    weight,
+    reduction,
+    ignore_index,
+    label_smoothing,
+):
+    # none reduction : onnx::Constant[value={0}]
+    # mean reduction : onnx::Constant[value={1}]
+    # sum reduction : onnx::Constant[value={2}]
+    reduction = symbolic_helper._maybe_get_const(reduction, "i")
+    reduction_vals = ["none", "mean", "sum"]
+    reduction = reduction_vals[reduction]
+
+    label_smoothing = symbolic_helper._maybe_get_const(label_smoothing, "f")
+    if label_smoothing is not None and label_smoothing > 0.0:
+        raise errors.SymbolicValueError(
+            "Unsupported: ONNX does not support label_smoothing", self
+        )
+
+    # in onnx SoftmaxCrossEntropyLoss specification, ignore_index is optional without default value.
+    # therefore we need to set ignore_index attribute even if it is not specified (e.g. ignore_index=-100).
+    ignore_index = symbolic_helper._maybe_get_const(ignore_index, "i")
+    if weight.node().mustBeNone():
+        celoss = g.op(
+            "SoftmaxCrossEntropyLoss",
+            self,
+            target,
+            reduction_s=reduction,
+            ignore_index_i=ignore_index,
+        )
+    else:
+        celoss = g.op(
+            "SoftmaxCrossEntropyLoss",
+            self,
+            target,
+            weight,
+            reduction_s=reduction,
+            ignore_index_i=ignore_index,
+        )
+
+    return celoss
+
+
+@_onnx_symbolic("aten::binary_cross_entropy_with_logits")
+@symbolic_helper.parse_args("v", "v", "v", "v", "i")
+def binary_cross_entropy_with_logits(
+    g: jit_utils.GraphContext, input, target, weight, pos_weight, reduction
+):
+    p = g.op("Constant", value_t=torch.tensor([1]))
+    sig_x = opset9.sigmoid(g, input)
+    log_sig_x = opset9.log(g, sig_x)
+    sub_1_x = opset9.sub(g, p, sig_x)
+    sub_1_y = opset9.sub(g, p, target)
+    log_1_x = opset9.log(g, sub_1_x)
+    if pos_weight is None or symbolic_helper._is_none(pos_weight):
+        output = opset9.neg(
+            g,
+            opset9.add(
+                g, opset9.mul(g, target, log_sig_x), opset9.mul(g, sub_1_y, log_1_x)
+            ),
+        )
+    else:
+        output = opset9.neg(
+            g,
+            opset9.add(
+                g,
+                opset9.mul(g, opset9.mul(g, target, log_sig_x), pos_weight),
+                opset9.mul(g, sub_1_y, log_1_x),
+            ),
+        )
+
+    if weight is not None and not symbolic_helper._is_none(weight):
+        output = opset9.mul(g, weight, output)
+
+    reduction = symbolic_helper._maybe_get_const(reduction, "i")
+    if reduction == 0:
+        return output
+    elif reduction == 1:
+        return g.op("ReduceMean", output, keepdims_i=0)
+    elif reduction == 2:
+        return g.op("ReduceSum", output, keepdims_i=0)
+    else:
+        return symbolic_helper._onnx_unsupported(
+            "binary_cross_entropy_with_logits with reduction other than none, mean, or sum",
+            input,
+        )
+
+
+@_onnx_symbolic("aten::celu")
+def celu(g: jit_utils.GraphContext, self, alpha):
+    alpha = symbolic_helper._maybe_get_const(alpha, "f")
+    # if the input is of type double cast it to float
+    if (
+        _type_utils.JitScalarType.from_value(self, _type_utils.JitScalarType.UNDEFINED)
+        == _type_utils.JitScalarType.DOUBLE
+    ):
+        self = g.op("Cast", self, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+        out = g.op("Celu", self, alpha_f=alpha)
+        return g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.DOUBLE)
+
+    return g.op("Celu", self, alpha_f=alpha)
+
+
+@_onnx_symbolic("aten::argmax")
+@symbolic_helper.parse_args("v", "v", "b")
+def argmax(
+    g: jit_utils.GraphContext,
+    input: torch._C.Value,
+    dim: torch._C.Value,
+    keepdim: bool,
+):
+    return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMax")
+
+
+@_onnx_symbolic("aten::argmin")
+@symbolic_helper.parse_args("v", "v", "b")
+def argmin(
+    g: jit_utils.GraphContext,
+    input: torch._C.Value,
+    dim: torch._C.Value,
+    keepdim: bool,
+):
+    return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMin")
+
+
+@_onnx_symbolic("aten::pow")
+def pow(g: jit_utils.GraphContext, self, exponent):
+    return g.op("Pow", self, exponent)
+
+
+@_onnx_symbolic("aten::ge")
+def ge(g: jit_utils.GraphContext, input, other):
+    return g.op("GreaterOrEqual", input, other)
+
+
+@_onnx_symbolic("aten::le")
+def le(g: jit_utils.GraphContext, input, other):
+    return g.op("LessOrEqual", input, other)
+
+
+@_onnx_symbolic("aten::unfold")
+@symbolic_helper.parse_args("v", "i", "v", "v")
+def unfold(g: jit_utils.GraphContext, input, dimension, size, step):
+    const_size = symbolic_helper._maybe_get_const(size, "i")
+    const_step = symbolic_helper._maybe_get_const(step, "i")
+    if not symbolic_helper._is_value(const_size) and not symbolic_helper._is_value(
+        const_step
+    ):
+        return opset9.unfold(g, input, dimension, const_size, const_step)
+
+    sizedim = symbolic_helper._get_tensor_dim_size(input, dimension)
+    if sizedim is not None:
+        low_start = g.op("Constant", value_t=torch.tensor(0))
+        low_end = g.op("Constant", value_t=torch.tensor(sizedim))
+        hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1))
+        low_indices = g.op("Range", low_start, low_end, step)
+        hi_indices = g.op("Range", size, hi_end, step)
+
+        low_size = symbolic_helper._size_helper(
+            g, low_indices, g.op("Constant", value_t=torch.tensor(0))
+        )
+        hi_size = symbolic_helper._size_helper(
+            g, hi_indices, g.op("Constant", value_t=torch.tensor(0))
+        )
+
+        ndim = symbolic_helper._get_tensor_rank(input)
+        if ndim is None:
+            raise AssertionError("ndim must be non-None")
+        perm = list(range(ndim))
+        perm.append(perm.pop(dimension))
+
+        unsqueeze_list = []
+        loop_condition = g.op("Constant", value_t=torch.tensor(1))
+        loop_condition = g.op(
+            "Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL
+        )
+        loop_len = g.op("Min", low_size, hi_size)
+
+        loop, (loop_context,), _ = jit_utils.add_op_with_blocks(
+            g, "Loop", loop_len, loop_condition, n_blocks=1
+        )
+
+        loop_block = loop_context.block
+        block_input_iter = utils._add_input_to_block(loop_block)
+        cond = utils._add_input_to_block(loop_block)  # noqa: F841
+
+        starts = loop_context.op("Gather", low_indices, block_input_iter)
+        ends = loop_context.op("Gather", hi_indices, block_input_iter)
+        axes = loop_context.op("Constant", value_t=torch.tensor([2]))
+        starts = symbolic_helper._unsqueeze_helper(loop_context, starts, [0])
+        ends = symbolic_helper._unsqueeze_helper(loop_context, ends, [0])
+        stack = loop_context.op("Slice", input, starts, ends, axes)
+
+        unsqueeze = symbolic_helper._unsqueeze_helper(
+            loop_context, loop_context.op("Transpose", stack, perm_i=perm), [dimension]
+        )
+        unsqueeze_list.append(unsqueeze)
+        concat = loop_context.op("Concat", *unsqueeze_list, axis_i=0)
+
+        cond_out = loop_context.op(
+            "Cast",
+            loop_condition,
+            # pyrefly: ignore [bad-argument-type]
+            _C_onnx.TensorProtoDataType.BOOL,
+        )
+        utils._add_output_to_block(loop_block, cond_out)
+        utils._add_output_to_block(loop_block, concat)
+
+        loop_output = loop.node().output()
+        perm = [0, 1, 2, 3, 4]
+        perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0]
+        transpose = g.op("Transpose", loop_output, perm_i=perm)
+        squeeze = symbolic_helper._squeeze_helper(g, transpose, [0])
+
+        return squeeze
+
+    return symbolic_helper._unimplemented("Unfold", "input size not accessible")
+
+
+@_onnx_symbolic("aten::tensordot")
+@symbolic_helper.parse_args("v", "v", "is", "is", "v")
+def tensordot(g: jit_utils.GraphContext, input_a, input_b, dims_a, dims_b, out=None):
+    if out is not None:
+        symbolic_helper._unimplemented(
+            "Tensordot", "Out parameter is not supported for tensordot."
+        )
+
+    dim_count_a = symbolic_helper._get_tensor_rank(input_a)
+    if dim_count_a is None:
+        raise errors.SymbolicValueError(
+            "Unsupported: ONNX export of tensordot for tensor(input_a) of unknown rank.",
+            input_a,
+        )
+
+    dim_count_b = symbolic_helper._get_tensor_rank(input_b)
+    if dim_count_b is None:
+        raise errors.SymbolicValueError(
+            "Unsupported: ONNX export of tensordot for tensor(input_b) of unknown rank.",
+            input_b,
+        )
+
+    dims_a = [
+        (dims_a[i] + dim_count_a) if (dims_a[i] < 0) else dims_a[i]
+        for i in range(len(dims_a))
+    ]
+    dims_b = [
+        (dims_b[i] + dim_count_b) if (dims_b[i] < 0) else dims_b[i]
+        for i in range(len(dims_b))
+    ]
+
+    left_dims_a = [i for i in range(dim_count_a) if (i not in dims_a)]
+    left_dims_b = [i for i in range(dim_count_b) if (i not in dims_b)]
+
+    new_input_a = opset9.permute(g, input_a, left_dims_a + dims_a)
+    new_input_b = opset9.permute(g, input_b, dims_b + left_dims_b)
+
+    input_shape = g.op("Shape", new_input_a)
+    left_sizes_a = symbolic_helper._slice_helper(
+        g, input_shape, axes=[0], starts=[0], ends=[len(left_dims_a)]
+    )
+    shape_sizes = [
+        left_sizes_a,
+        g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
+    ]
+    output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes)
+
+    input_shape = g.op("Shape", output_a)
+    slices = symbolic_helper._slice_helper(
+        g, input_shape, axes=[0], starts=[-1], ends=[sys.maxsize]
+    )
+    shape_sizes = [
+        g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
+        slices,
+    ]
+    output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes)
+
+    input_shape = g.op("Shape", new_input_b)
+    left_sizes_b = symbolic_helper._slice_helper(
+        g, input_shape, axes=[0], starts=[len(dims_b)], ends=[sys.maxsize]
+    )
+    slices = symbolic_helper._slice_helper(
+        g, input_shape, axes=[0], starts=[0], ends=[len(dims_b)]
+    )
+    shape_sizes = [
+        slices,
+        g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
+    ]
+    output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes)
+
+    input_shape = g.op("Shape", output_b)
+    slices = symbolic_helper._slice_helper(
+        g, input_shape, axes=[0], starts=[-1], ends=[sys.maxsize]
+    )
+    shape_sizes = [
+        g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
+        slices,
+    ]
+    output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes)
+
+    output = einsum(g, "ij,jk->ik", g.op("prim::ListConstruct", *[output_a, output_b]))
+
+    shape_sizes = [left_sizes_a, left_sizes_b]
+    return opset9._reshape_from_tensor(g, output, shape_sizes)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py

@@ -0,0 +1,301 @@
+# mypy: allow-untyped-defs
+# mypy: disable-error-code=arg-type
+"""This file exports ONNX ops for opset 14.
+
+Note [ONNX operators that are added/updated in opset 14]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+New operators:
+    HardSwish, Trilu
+
+Updated operators:
+    Reshape
+    Add, Sub, Mul, Div
+    GRU, LSTM, RNN
+    BatchNorm, Cumsum, Relu
+"""
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in README.md
+from __future__ import annotations
+
+import functools
+
+import torch
+from torch.onnx import _constants
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+)
+from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
+
+
+__all__ = [
+    "hardswish",
+    "tril",
+    "triu",
+    "reshape",
+    "batch_norm",
+    "quantized_hardswish",
+    "scaled_dot_product_attention",
+]
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=14)
+
+
+@_onnx_symbolic("aten::hardswish")
+@symbolic_helper.parse_args("v")
+def hardswish(g: jit_utils.GraphContext, self):
+    return g.op("HardSwish", self)
+
+
+@_onnx_symbolic("aten::tril")
+def tril(g: jit_utils.GraphContext, self, diagonal, out=None):
+    return g.op("Trilu", self, diagonal, upper_i=0)
+
+
+@_onnx_symbolic("aten::triu")
+def triu(g: jit_utils.GraphContext, self, diagonal, out=None):
+    return g.op("Trilu", self, diagonal, upper_i=1)
+
+
+@_onnx_symbolic("aten::reshape")
+@symbolic_helper.quantized_args(True)
+@symbolic_helper.parse_args("v", "v")
+def reshape(g: jit_utils.GraphContext, self, shape):
+    # NOTE: Due to bug in ORT https://github.com/microsoft/onnxruntime/issues/10664
+    #       Reshape export cannot utilize the new allowzero attribute introduced in opset 14.
+    return symbolic_helper._reshape_helper(g, self, shape, allowzero=0)
+
+
+@_onnx_symbolic("aten::batch_norm")
+@symbolic_helper.parse_args("v", "v", "v", "v", "v", "i", "f", "f", "i")
+def batch_norm(
+    g: jit_utils.GraphContext,
+    input,
+    weight,
+    bias,
+    running_mean,
+    running_var,
+    training,
+    momentum,
+    eps,
+    cudnn_enabled,
+):
+    if (
+        torch.is_autocast_enabled()
+        and not symbolic_helper.args_have_same_dtype(
+            [input, weight, bias, running_mean, running_var]
+        )
+        and GLOBALS.export_onnx_opset_version < 15
+    ):
+        return symbolic_helper._onnx_opset_unsupported_detailed(
+            "BatchNormalization",
+            14,
+            15,
+            "All input tensors must have the same `dtype`."
+            " Turn off Autocast or export using opset version 15.",
+            input,
+        )
+
+    symbolic_helper.check_training_mode(training, "batch_norm")
+    weight, bias, running_mean, running_var = symbolic_helper._batchnorm_helper(
+        g, input, weight, bias, running_mean, running_var
+    )
+    out = g.op(
+        "BatchNormalization",
+        input,
+        weight,
+        bias,
+        running_mean,
+        running_var,
+        epsilon_f=eps,
+        momentum_f=1 - momentum,
+        training_mode_i=0 if not training else 1,
+        outputs=1 if not training else 3,
+    )
+    if not training:
+        return out
+    else:
+        res, new_running_mean, new_running_var = out
+        new_running_mean.setType(running_mean.type())
+        new_running_var.setType(running_var.type())
+        return res
+
+
+@_onnx_symbolic("quantized::hardswish")
+def quantized_hardswish(g: jit_utils.GraphContext, x, op_scale, op_zero_point):
+    x, _, _, _ = symbolic_helper.dequantize_helper(g, x)
+
+    output = hardswish(g, x)
+
+    return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point)
+
+
+# Ported from
+# https://github.com/microsoft/onnxscript/blob/6b1b81700b4523f31d8c6d3321e5d8ef5d42b764/onnxscript/function_libs/torch_aten/ops/nn.py#L1504
+# aten_scaled_dot_product_attention
+# NOTE: Need op.Trilu
+@_onnx_symbolic("aten::scaled_dot_product_attention")
+@symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b")
+def scaled_dot_product_attention(
+    g: jit_utils.GraphContext,
+    query: torch._C.Value,
+    key: torch._C.Value,
+    value: torch._C.Value,
+    attn_mask: torch._C.Value | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: torch._C.Value | None = None,
+    enable_gqa: bool = False,
+):
+    if is_causal and not symbolic_helper._is_none(attn_mask):
+        raise AssertionError("is_causal and attn_mask cannot be set at the same time")
+    if enable_gqa:
+        raise AssertionError(
+            "conversion of scaled_dot_product_attention not implemented if enable_gqa is True"
+        )
+
+    if symbolic_helper._is_none(scale):
+        scale = _attention_scale(g, query)
+
+    if is_causal:
+        attn_mask = _causal_attention_mask(g, query, key)
+
+    # Swap the last two axes of key
+    # NOTE: onnx-script has different logic here, because the attribute perms in
+    # transpose needs list of ints
+    key_shape_builtin = symbolic_helper._get_tensor_rank(key)
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
+    key_transposed_axes = list(range(key_shape_builtin))
+    key_transposed_axes[-1], key_transposed_axes[-2] = (
+        key_transposed_axes[-2],
+        key_transposed_axes[-1],
+    )
+    key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    # pyrefly: ignore [bad-argument-type]
+    query_scaled = g.op("Mul", query, g.op("Sqrt", scale))
+    # pyrefly: ignore [bad-argument-type]
+    key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale))
+    mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled)
+
+    if symbolic_helper._is_none(attn_mask):
+        mul_qk_add = mul_qk
+        attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1)
+    elif (
+        _type_utils.JitScalarType.from_value(attn_mask)
+        == _type_utils.JitScalarType.BOOL
+    ):
+        # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf'))
+        const_zero = g.op("Constant", value_t=torch.tensor([0.0]))
+        const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")]))
+        # pyrefly: ignore [bad-argument-type]
+        attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf)
+        mul_qk_add = g.op("Add", mul_qk, attn_mask)
+        attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1)
+        # When using scaled dot product attention with a boolean mask, the softmax operation might return NaN values
+        # due to the presence of -inf in an entire row (padding tokens), resulting in 0/0 (NaN) in the softmax output.
+        # This is because there's no safe softmax imp in ONNX, so we need to handle NaN values explicitly to match
+        # the behavior of PyTorch with boolean masks.
+        attn_weight = g.op("Where", g.op("IsNaN", attn_weight), const_zero, attn_weight)
+    elif _type_utils.JitScalarType.from_value(attn_mask) in (
+        _type_utils.JitScalarType.FLOAT,
+        _type_utils.JitScalarType.HALF,
+        _type_utils.JitScalarType.BFLOAT16,
+    ):
+        # pyrefly: ignore [bad-argument-type]
+        mul_qk_add = g.op("Add", mul_qk, attn_mask)
+        attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1)
+    else:
+        raise ValueError(
+            f"Unsupported type for attn_mask: {_type_utils.JitScalarType.from_value(attn_mask)}"
+        )
+
+    if dropout_p != 0:
+        attn_weight = g.op(
+            "Dropout",
+            attn_weight,
+            g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)),
+        )
+
+    return g.op("MatMul", attn_weight, value)
+
+
+def _attention_scale(
+    g: jit_utils.GraphContext, query: torch._C.Value
+) -> torch._C.Value:
+    """Calculate the scale factor for the attention result.
+
+    Args:
+        query: Tensor of shape [..., L, E]
+
+    Returns:
+        Scalar scale factor := 1 / math.sqrt(query.size(-1))
+    """
+    query_shape = g.op("Shape", query)
+    query_shape_last = g.op(
+        "Slice",
+        query_shape,
+        g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)),
+        g.op(
+            "Constant", value_t=torch.tensor([_constants.INT64_MAX], dtype=torch.int64)
+        ),
+    )
+    embedding_size = g.op(
+        "Cast",
+        query_shape_last,
+        to_i=_type_utils.JitScalarType.from_value(query).onnx_type(),
+    )
+    const_one = g.op("Constant", value_t=torch.tensor([1.0], dtype=torch.float))
+    scale = g.op("Div", const_one, g.op("Sqrt", embedding_size))
+    # Add a Cast to convert the scale back to original type
+    scale = g.op(
+        "Cast",
+        scale,
+        to_i=_type_utils.JitScalarType.from_value(query).onnx_type(),
+    )
+    return scale
+
+
+def _causal_attention_mask(
+    g: jit_utils.GraphContext, query: torch._C.Value, key: torch._C.Value
+) -> torch._C.Value:
+    """Create a causal mask for the given query and key tensors.
+
+    Equivalent to::
+        mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_mask = torch.zeros(L, S, dtype=torch.float)
+        attn_mask = attn_mask.masked_fill(not mask, -float("inf"))
+
+    Args:
+        query: Tensor of shape [..., L, E]
+        key: Tensor of shape [..., S, E]
+
+    Returns:
+        Tensor of shape [L, S]
+    """
+
+    query_shape = g.op("Shape", query)
+    key_shape = g.op("Shape", key)
+
+    last_idx = g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64))
+    second_last_idx = g.op("Constant", value_t=torch.tensor([-2], dtype=torch.int64))
+    target_length = g.op("Slice", query_shape, second_last_idx, last_idx)
+    source_length = g.op("Slice", key_shape, second_last_idx, last_idx)
+    # attn_mask = torch.ones(L, S) := {
+    size = g.op("Concat", target_length, source_length, axis_i=0)
+    const_one = g.op("Constant", value_t=torch.tensor([1.0]))
+    attn_mask = g.op("Expand", const_one, size)
+    # }
+    attn_mask = g.op("Trilu", attn_mask, upper_i=0)
+    # The causal mask has 0s in the lower triangle and -inf in the upper triangle.
+    const_zero = g.op("Constant", value_t=torch.tensor([0.0]))
+    const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")]))
+    attn_mask = g.op(
+        "Where", g.op("Equal", attn_mask, const_zero), const_neg_inf, const_zero
+    )
+    return attn_mask

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset15.py

@@ -0,0 +1,84 @@
+# mypy: allow-untyped-defs
+"""This file exports ONNX ops for opset 15.
+
+Note [ONNX operators that are added/updated in opset 15]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/master/docs/Changelog.md#version-15-of-the-default-onnx-operator-set
+New operators:
+    Bernoulli
+    CastLike
+    Optional
+    OptionalGetElement
+    OptionalHasElement
+
+Updated operators:
+    BatchNormalization https://github.com/onnx/onnx/pull/3545
+                        Backwards compatible
+                        TODO: test coverage for mixed types inputs.
+    Pow                https://github.com/onnx/onnx/pull/3412
+                        Backwards compatible
+                        TODO: bfloat16 support.
+    Shape              https://github.com/onnx/onnx/pull/3580
+                        Backwards compatible
+                        TODO: optional start/end attribute.
+"""
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in README.md
+
+import functools
+
+import torch
+from torch import _C
+from torch.onnx._internal.torchscript_exporter import (
+    jit_utils,
+    registration,
+    symbolic_helper,
+    symbolic_opset9 as opset9,
+)
+
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=15)
+
+
+@_onnx_symbolic("aten::__is_")
+def aten__is_(g: jit_utils.GraphContext, self, other):
+    if symbolic_helper._is_none(other):
+        if isinstance(self.type(), _C.OptionalType):
+            none = g.op("OptionalHasElement", self)
+            return g.op("Not", none)
+        else:
+            return g.op("Constant", value_t=torch.BoolTensor([0]))
+    return opset9.eq(g, self, other)
+
+
+@_onnx_symbolic("aten::__isnot_")
+@opset9.wrap_logical_op_with_negation  # type: ignore[has-type]
+def aten__isnot_(g: jit_utils.GraphContext, self, other):
+    return aten__is_(g, self, other)
+
+
+@_onnx_symbolic("aten::bernoulli")
+def bernoulli(g: jit_utils.GraphContext, input, p=None, generator=None, out=None):
+    if out is not None and not symbolic_helper._is_none(out):
+        symbolic_helper._unimplemented(
+            "Bernoulli", "out parameter is not supported for bernoulli", input
+        )
+    if generator is not None and not symbolic_helper._is_none(generator):
+        symbolic_helper._unimplemented(
+            "Bernoulli", "generator is not supported for bernoulli", input
+        )
+    if p is None or symbolic_helper._is_none(p):
+        return g.op("Bernoulli", input)
+    return opset9.bernoulli(g, input, p, generator, out)
+
+
+@_onnx_symbolic("prim::unchecked_cast")
+def prim_unchecked_cast(g: jit_utils.GraphContext, self):
+    # exists to refine the type of the Value
+    # if x is Optional[Tensor], unchecked_cast will cast
+    # x to Tensor, so the rest of the graph knows that x is a Tensor.
+    if isinstance(self.type(), _C.OptionalType):
+        return g.op("OptionalGetElement", self)
+
+    return self

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset16.py

@@ -0,0 +1,191 @@
+# mypy: allow-untyped-defs
+"""This file exports ONNX ops for opset 16.
+
+Note [ONNX Operators that are added/updated in opset 16]
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-16-of-the-default-onnx-operator-set
+New operators:
+    GridSample https://github.com/onnx/onnx/pull/3557
+
+Updated operators:
+    Identity
+    If
+    LeakyRelu
+    Loop
+    PRelu
+    RoiAlign
+    Scan
+    ScatterElements
+    ScatterND
+    Where
+    GreaterOrEqual
+    LessOrEqual
+"""
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in README.md
+
+import functools
+
+import torch
+from torch.nn.functional import (
+    GRID_SAMPLE_INTERPOLATION_MODES,
+    GRID_SAMPLE_PADDING_MODES,
+)
+from torch.onnx import errors
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+    utils,
+)
+
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=16)
+
+
+# note (mkozuki): Why `grid_sampler` instead of `grid_sample`?
+# Because `torch.nn.functional.grid_sample` calls `torch.grid_sampler`.
+@_onnx_symbolic("aten::grid_sampler")
+@symbolic_helper.parse_args("v", "v", "i", "i", "b")
+def grid_sampler(
+    g: jit_utils.GraphContext,
+    input,
+    grid,
+    mode_enum,
+    padding_mode_enum,
+    align_corners,
+):
+    # Check the input and grid tensor rank beforehand.
+    if symbolic_helper._get_tensor_rank(input) == 5:
+        return symbolic_helper._onnx_unsupported("GridSample with 5D volumetric input")
+    mode_s = {v: k for k, v in GRID_SAMPLE_INTERPOLATION_MODES.items()}[mode_enum]  # type: ignore[call-arg]
+    padding_mode_s = {v: k for k, v in GRID_SAMPLE_PADDING_MODES.items()}[  # type: ignore[call-arg]
+        padding_mode_enum
+    ]
+    return g.op(
+        "GridSample",
+        input,
+        grid,
+        align_corners_i=int(align_corners),
+        mode_s=mode_s,
+        padding_mode_s=padding_mode_s,
+    )
+
+
+@_onnx_symbolic("aten::scatter_add")
+@symbolic_helper.parse_args("v", "i", "v", "v")
+def scatter_add(g: jit_utils.GraphContext, self, dim, index, src):
+    src_type = _type_utils.JitScalarType.from_value(
+        src, _type_utils.JitScalarType.UNDEFINED
+    )
+    src_sizes = symbolic_helper._get_tensor_sizes(src)
+    index_sizes = symbolic_helper._get_tensor_sizes(index)
+
+    if len(src_sizes) != len(index_sizes):
+        return symbolic_helper._unimplemented(
+            "scatter_add",
+            f"`index` ({index_sizes}) should have the same dimensionality as `src` ({src_sizes})",
+        )
+
+    # PyTorch only allows index shape <= src shape, so we can only consider
+    # taking index as subset size to src, like PyTorch does. When sizes for src
+    # and index are not matched or there are dynamic axes, we take index shape to
+    # slice src to accommodate.
+    if src_sizes != index_sizes or None in index_sizes:
+        adjusted_shape = g.op("Shape", index)
+        starts = g.op("Constant", value_t=torch.tensor([0] * len(index_sizes)))
+        src = g.op("Slice", src, starts, adjusted_shape)
+
+    src = symbolic_helper._maybe_get_scalar(src)
+    if symbolic_helper._is_value(src):
+        return g.op("ScatterElements", self, index, src, axis_i=dim, reduction_s="add")
+    else:
+        # Check if scalar "src" has same type as self (PyTorch allows different
+        # type for scalar src (but not when src is tensor)). If not, insert Cast node.
+        if _type_utils.JitScalarType.from_value(self) != src_type:
+            src = g.op(
+                "Cast",
+                src,
+                to_i=_type_utils.JitScalarType.from_value(self).onnx_type(),
+            )
+
+        return g.op(
+            "ScatterElements",
+            self,
+            index,
+            src,
+            axis_i=dim,
+            reduction_s="add",
+        )
+
+
+@_onnx_symbolic("aten::scatter_reduce")
+@symbolic_helper.parse_args("v", "i", "v", "v", "s", "b")
+def scatter_reduce(
+    g: jit_utils.GraphContext,
+    self: torch._C.Value,
+    dim: int,
+    index: torch._C.Value,
+    src: torch._C.Value,
+    reduce: str,
+    include_self: bool,
+):
+    if reduce == "mean":
+        raise errors.OnnxExporterError(
+            "ONNX does not support mean reduction for scatter_reduce"
+        )
+    if not include_self:
+        raise errors.OnnxExporterError(
+            "ONNX does not support include_self=False for scatter_reduce"
+        )
+
+    reduce_mode = {  # convert torch string name to onnx string name
+        "mean": "none",  # 'mean' doesn't support in ONNX 1.14 definition
+        "sum": "add",
+        "prod": "mul",
+        "amin": "min",
+        "amax": "max",
+    }
+    onnx_reduce = reduce_mode[reduce]
+
+    self_rank = g.op("Size", g.op("Shape", self))
+
+    # if self_rank == 0:  # assert (index_rank == 0 and rank_src == 0)
+    self_rank_is_zero = g.op(
+        "Equal", self_rank, g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64))
+    )
+    if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks(
+        g, "If", self_rank_is_zero, n_blocks=2, outputs=3
+    )
+    neg_1 = if_context.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64))
+
+    self_reshape = if_context.op("Reshape", self, neg_1)
+    utils._add_output_to_block(if_context.block, self_reshape)
+    index_reshape = if_context.op("Reshape", index, neg_1)
+    utils._add_output_to_block(if_context.block, index_reshape)
+    src_reshape = if_context.op("Reshape", src, neg_1)
+    utils._add_output_to_block(if_context.block, src_reshape)
+
+    self_identity = else_context.op("Identity", self)
+    utils._add_output_to_block(else_context.block, self_identity)
+    index_identitye = else_context.op("Identity", index)
+    utils._add_output_to_block(else_context.block, index_identitye)
+    src_identity = else_context.op("Identity", src)
+    utils._add_output_to_block(else_context.block, src_identity)
+
+    result = g.op("ScatterElements", *if_op, axis_i=dim, reduction_s=onnx_reduce)
+
+    # if self_rank == 0:
+    if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks(
+        g, "If", self_rank_is_zero, n_blocks=2, outputs=1
+    )
+    result_squeezed = if_context.op("Squeeze", result)
+    utils._add_output_to_block(if_context.block, result_squeezed)
+    result_identity = else_context.op("Identity", result)
+    utils._add_output_to_block(else_context.block, result_identity)
+    result_final = if_op.node().output()
+
+    return result_final

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py

@@ -0,0 +1,252 @@
+# mypy: allow-untyped-defs
+# mypy: disable-error-code=arg-type
+"""This file exports ONNX ops for opset 17.
+
+Note [ONNX Operators that are added/updated in opset 17]
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-17-of-the-default-onnx-operator-set
+New operators:
+    BlackmanWindow
+    DFT
+    HammingWindow
+    HannWindow
+    LayerNormalization
+    MelWeightMatrix
+    STFT
+    SequenceMap
+"""
+
+import functools
+from collections.abc import Sequence
+
+import torch
+from torch import _C
+from torch.onnx import errors
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+)
+
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in README.md
+
+__all__ = ["layer_norm", "stft", "quantized_layer_norm"]
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=17)
+
+
+@_onnx_symbolic("aten::layer_norm")
+@symbolic_helper.parse_args("v", "is", "v", "v", "f", "none")
+def layer_norm(
+    g: jit_utils.GraphContext,
+    input: _C.Value,
+    normalized_shape: Sequence[int],
+    weight: _C.Value,
+    bias: _C.Value,
+    eps: float,
+    cudnn_enable: bool,
+):
+    # normalized_shape: input shape from an expected input of size
+    # axis: The first normalization dimension.
+    # layer_norm normalizes on the last D dimensions,
+    # where D is the size of normalized_shape
+    axis = -len(normalized_shape)
+    scalar_type = _type_utils.JitScalarType.from_value(
+        input, _type_utils.JitScalarType.FLOAT
+    )
+    dtype = scalar_type.dtype()
+    if symbolic_helper._is_none(weight):
+        weight_value = torch.ones(normalized_shape, dtype=dtype)
+        weight = g.op("Constant", value_t=weight_value)
+    if symbolic_helper._is_none(bias):
+        bias_value = torch.zeros(normalized_shape, dtype=dtype)
+        bias = g.op("Constant", value_t=bias_value)
+    return g.op(
+        "LayerNormalization",
+        input,
+        weight,
+        bias,
+        epsilon_f=eps,
+        axis_i=axis,
+    )
+
+
+@_onnx_symbolic("quantized::layer_norm")
+def quantized_layer_norm(
+    g: jit_utils.GraphContext,
+    x,
+    normalized_shape,
+    weight,
+    bias,
+    eps,
+    op_scale,
+    op_zero_point,
+):
+    x, _, _, _ = symbolic_helper.dequantize_helper(g, x)
+
+    output = layer_norm(g, x, normalized_shape, weight, bias, eps, False)
+
+    return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point)
+
+
+def _compute_edge_sizes(n_fft, window_size):
+    """Helper function to compute the sizes of the edges (left and right)
+    of a given window centered within an FFT size."""
+    left = (n_fft - window_size) // 2
+    right = n_fft - left - window_size
+    return left, right
+
+
+@_onnx_symbolic("aten::stft")
+@symbolic_helper.parse_args("v", "i", "i", "i", "v", "b", "b", "b", "b")
+def stft(
+    g: jit_utils.GraphContext,
+    input: _C.Value,
+    n_fft: int,
+    hop_length: int | None = None,
+    win_length: int | None = None,
+    window: _C.Value | None = None,
+    normalized: bool = False,
+    onesided: bool | None = True,
+    return_complex: bool | None = False,
+    align_to_window: bool | None = None,
+) -> _C.Value:
+    """Associates `torch.stft` with the `STFT` ONNX operator.
+    Note that torch.stft calls _VF.stft, without centering or padding options.
+    Hence, this function does not contain these two arguments.
+    See torch.stft source code for more info.
+
+    Args:
+        g: Graph to write the ONNX representation into
+        input: Input tensor for the transformation
+        n_fft: FFT size
+        hop_length: Size of the hop. Defaults to `floot(n_fft // 4)`
+        win_length: Size of the analysis window. Defaults to `n_fft`
+        window: Analysis window. Defaults to a window of all ones
+        normalized: Whether to return a normalized STFT
+        onesided: Whether to return only half (+1) of the results, given the
+            symmetry of the STFT
+        return_complex: Whether to return the complex value (Note: Must be
+            `False` or `None`)
+
+    Returns:
+        op: Operator for torch.stft associated with STFT (ONNX)
+    """
+    # Checks
+    if return_complex:
+        raise errors.SymbolicValueError(
+            msg="STFT does not currently support complex types", value=input
+        )
+
+    if align_to_window is not None:
+        raise errors.SymbolicValueError(
+            msg="STFT does not currently support the align_to_window option",
+            value=input,
+        )  # TODO(#145944): add compatibility with align_to_window option.
+
+    # Get STFT sizes
+    frame_step_value = hop_length if hop_length is not None else n_fft // 4
+    frame_step_const = g.op(
+        "Constant", value_t=torch.tensor(frame_step_value, dtype=torch.int64)
+    )
+    frame_length_const = g.op(
+        "Constant", value_t=torch.tensor(n_fft, dtype=torch.int64)
+    )
+
+    # Pre-process input if needed
+    signal = input
+    signal_rank = symbolic_helper._get_tensor_rank(signal)
+    if signal_rank == 1:
+        # Add batch dimension
+        signal = g.op(
+            "Unsqueeze",
+            signal,
+            g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)),
+        )
+    elif signal_rank is None or signal_rank > 2:
+        raise errors.SymbolicValueError(
+            msg="STFT can only take inputs of 1 [signal] or 2 [batch, signal] dimensions. "
+            f"Current rank of signal is {signal_rank}, please reduce it.",
+            value=input,
+        )
+
+    # Get window and make sure it's the same size as `win_length` or `n_fft`
+    # pyrefly: ignore [bad-argument-type]
+    n_win = symbolic_helper._get_tensor_dim_size(window, dim=0)
+    if n_win is not None:
+        win_length_default = win_length if win_length else n_fft
+        if n_win != win_length_default:
+            raise AssertionError(
+                "Analysis window size must equal `win_length` or `n_fft`. "
+                f"Please, set `win_length` or `n_fft` to match `window` size ({n_win})"
+            )
+
+        # Center window around zeros if needed (required by ONNX's STFT)
+        if n_win < n_fft:
+            left, right = _compute_edge_sizes(n_fft, n_win)
+            left_win = g.op("Constant", value_t=torch.zeros(left))
+            right_win = g.op("Constant", value_t=torch.zeros(right))
+            # pyrefly: ignore [bad-argument-type]
+            window = g.op("Concat", left_win, window, right_win, axis_i=0)
+
+    # Create window, if needed
+    if symbolic_helper._is_none(window):
+        if win_length:
+            if win_length > n_fft:
+                raise errors.SymbolicValueError(
+                    msg="The analysis window can't be longer than the size of the FFT. "
+                    f"Please set `win_length` ({win_length}) to `n_fft` ({n_fft}) or less.",
+                    value=input,
+                )
+
+            # Center window, if needed
+            left, right = _compute_edge_sizes(n_fft, win_length)
+            torch_window = torch.hstack(
+                (torch.zeros(left), torch.ones(win_length), torch.zeros(right))
+            )
+        else:
+            # Rectangle window
+            torch_window = torch.ones(n_fft)
+        if torch_window.shape[0] != n_fft:
+            raise AssertionError(
+                f"torch_window.shape[0]={torch_window.shape[0]} != n_fft={n_fft}"
+            )
+        window = g.op("Constant", value_t=torch_window)
+    window = g.op(
+        "Cast",
+        # pyrefly: ignore [bad-argument-type]
+        window,
+        to_i=_type_utils.JitScalarType.from_value(signal).onnx_type(),
+    )
+
+    # Run STFT
+    result = g.op(
+        "STFT",
+        signal,
+        frame_step_const,
+        window,
+        frame_length_const,
+        onesided_i=1 if onesided is None or onesided else 0,
+    )
+
+    # Transpose to mimic torch.stft's behavior
+    result = g.op("Transpose", result, perm_i=[0, 2, 1, 3])
+
+    # Remove batch dimension, if needed
+    if signal_rank == 1:
+        result = g.op(
+            "Squeeze",
+            result,
+            g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)),
+        )
+
+    # Normalize, if needed
+    if normalized:
+        sqrt_nfft = torch.sqrt(torch.tensor(n_fft, dtype=signal.type().dtype()))
+        result = g.op("Div", result, g.op("Constant", value_t=sqrt_nfft))
+
+    return result

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py

@@ -0,0 +1,272 @@
+# mypy: allow-untyped-defs
+"""This file exports ONNX ops for opset 18.
+
+Note [ONNX Operators that are added/updated in opset 18]
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-18-of-the-default-onnx-operator-set
+New operators:
+    BitwiseAnd
+    CenterCropPad
+    Col2Im
+    Mish
+    OptionalGetElement
+    OptionalHasElement
+    Pad
+    Resize
+    ScatterElements
+    ScatterND
+    Split
+"""
+
+import functools
+from collections.abc import Sequence
+
+import torch
+from torch import _C
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+    symbolic_opset9 as opset9,
+)
+
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in symbolic_helper.py
+
+__all__ = [
+    "col2im",
+]
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=18)
+
+
+@_onnx_symbolic("aten::__and_")
+@_onnx_symbolic("aten::bitwise_and")
+def __and_(g: jit_utils.GraphContext, self, other):
+    # do type promotion (scalars don't seem to apply)
+    args = [self, other]
+    # type promotion doesn't happen with torch.bitwise_and(tensor, scalar)
+    prom_args = [arg for arg in args if symbolic_helper._get_tensor_rank(arg)]
+    if len(prom_args) == 0:
+        prom_args = args
+    promotion_jit_type = symbolic_helper._type_promote_from_values(*prom_args)
+    self = symbolic_helper._maybe_cast_to_type(g, self, promotion_jit_type)
+    other = symbolic_helper._maybe_cast_to_type(g, other, promotion_jit_type)
+    if promotion_jit_type == _type_utils.JitScalarType.BOOL:
+        return g.op("And", self, other)
+    return g.op("BitwiseAnd", self, other)
+
+
+@_onnx_symbolic("aten::col2im")
+@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is")
+def col2im(
+    g,
+    input: _C.Value,
+    output_size: _C.Value,
+    kernel_size: _C.Value,
+    dilation: Sequence[int],
+    padding: Sequence[int],
+    stride: Sequence[int],
+):
+    # convert [i0, i1, ..., in] into [i0, i0, i1, i1, ..., in, in]
+    adjusted_padding: list[int] = []
+    for pad in padding:
+        adjusted_padding.extend(pad for _ in range(2))
+
+    num_dimensional_axis = symbolic_helper._get_tensor_sizes(output_size)[0]
+    if not adjusted_padding:
+        adjusted_padding = [0, 0] * num_dimensional_axis
+
+    if not dilation:
+        dilation = [1] * num_dimensional_axis
+
+    if not stride:
+        stride = [1] * num_dimensional_axis
+
+    return g.op(
+        "Col2Im",
+        input,
+        output_size,
+        kernel_size,
+        dilations_i=dilation,
+        pads_i=adjusted_padding,
+        strides_i=stride,
+    )
+
+
+@_onnx_symbolic(
+    "aten::mean", decorate=[symbolic_helper._apply_params("ReduceMean", "mean")]
+)
+@_onnx_symbolic(
+    "aten::prod",
+    decorate=[
+        symbolic_helper._apply_params(
+            "ReduceProd", "prod", allow_multi_dim_support=False
+        )
+    ],
+)
+def _reduce_with_dtype(onnx_op: str, name: str, allow_multi_dim_support: bool = True):
+    return symbolic_helper._reduce_with_dtype_helper(
+        onnx_op, name, allow_multi_dim_support
+    )
+
+
+@_onnx_symbolic("aten::native_layer_norm")
+@symbolic_helper.quantized_args(True, False, False, False)
+@symbolic_helper.parse_args("v", "is", "v", "v", "f")
+def _native_layer_norm(
+    g: jit_utils.GraphContext,
+    input: _C.Value,
+    normalized_shape: Sequence[int],
+    weight: _C.Value,
+    bias: _C.Value,
+    eps: float,
+) -> tuple[_C.Value, _C.Value, _C.Value]:
+    return opset9.native_layer_norm(g, input, normalized_shape, weight, bias, eps)
+
+
+@_onnx_symbolic("aten::glu")
+@symbolic_helper.parse_args("v", "i")
+def _glu(g: jit_utils.GraphContext, input, dim):
+    dim_size = symbolic_helper._get_tensor_dim_size(input, dim)
+    if dim_size is not None:
+        if dim_size % 2 != 0:
+            raise AssertionError(f"dim_size must be even, got {dim_size}")
+
+    first, second = g.op("Split", input, axis_i=dim, num_outputs_i=2, outputs=2)
+    return g.op("Mul", first, g.op("Sigmoid", second))
+
+
+@_onnx_symbolic("aten::max")
+# torch.max (same for torch.min) actually has two interfaces smashed together:
+# torch.max(x, dim, keepdim) and torch.max(x, y)
+# TODO(justinchuby): Support multiple quantized args in output
+def max(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None):
+    return symbolic_helper._max_helper(g, self, dim_or_y, keepdim)
+
+
+@_onnx_symbolic("aten::maximum")
+@symbolic_helper.quantized_args(True, True)
+def maximum(g: jit_utils.GraphContext, input, other):
+    # pyrefly: ignore [no-matching-overload]
+    return max(g, input, dim_or_y=other)
+
+
+@_onnx_symbolic("aten::min")
+# TODO(justinchuby): Support multiple quantized args in output
+def min(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None):
+    return symbolic_helper._min_helper(g, self, dim_or_y, keepdim)
+
+
+@_onnx_symbolic("aten::minimum")
+@symbolic_helper.quantized_args(True, True)
+def minimum(g: jit_utils.GraphContext, input, other):
+    # pyrefly: ignore [no-matching-overload]
+    return min(g, input, dim_or_y=other)
+
+
+@_onnx_symbolic("aten::amax")
+@symbolic_helper.quantized_args(True)
+@symbolic_helper.parse_args("v", "is", "i")
+def amax(g: jit_utils.GraphContext, self, dim, keepdim):
+    axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long))
+    return g.op("ReduceMax", self, axes, keepdims_i=keepdim)
+
+
+@_onnx_symbolic("aten::amin")
+@symbolic_helper.quantized_args(True)
+@symbolic_helper.parse_args("v", "is", "i")
+def amin(g: jit_utils.GraphContext, self, dim, keepdim):
+    axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long))
+    return g.op("ReduceMin", self, axes, keepdims_i=keepdim)
+
+
+@_onnx_symbolic("aten::aminmax")
+@symbolic_helper.quantized_args(True)
+@symbolic_helper.parse_args("v", "v", "i")
+def aminmax(g: jit_utils.GraphContext, self, dim, keepdim):
+    if not symbolic_helper._is_none(dim):
+        dim = symbolic_helper._get_const(dim, "i", "dim")
+        axes = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long))
+        return g.op("ReduceMin", self, axes, keepdims_i=keepdim), g.op(
+            "ReduceMax", self, axes, keepdims_i=keepdim
+        )
+    else:
+        return g.op("ReduceMin", self, keepdims_i=keepdim), g.op(
+            "ReduceMax", self, keepdims_i=keepdim
+        )
+
+
+@_onnx_symbolic("aten::var_mean")
+def _var_mean(g: jit_utils.GraphContext, input, *args):
+    if len(args) == 1:
+        return symbolic_helper._var_mean_helper(g, input, None, args[0], None)
+    else:
+        return symbolic_helper._var_mean_helper(g, input, *args)
+
+
+@_onnx_symbolic("aten::logsumexp")
+@symbolic_helper.parse_args("v", "is", "i")
+def _logsumexp(g: jit_utils.GraphContext, input, dim, keepdim):
+    if dim is None:
+        return g.op("ReduceLogSumExp", input, keepdims_i=0)
+    else:
+        axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long))
+        return g.op("ReduceLogSumExp", input, axes, keepdims_i=keepdim)
+
+
+@_onnx_symbolic("aten::linalg_matrix_norm")
+@symbolic_helper.parse_args("v", "v", "is", "b", "v")
+def _linalg_matrix_norm(
+    g: jit_utils.GraphContext,
+    self: torch._C.Value,
+    ord: torch._C.Value,
+    dim: list[int],
+    keepdim: bool,
+    dtype: torch._C.Value,
+):
+    return opset9.linalg_matrix_norm(g, self, ord, dim, keepdim, dtype)
+
+
+@_onnx_symbolic("aten::embedding_bag")
+@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i", "v", "i", "i")
+def embedding_bag(
+    g: jit_utils.GraphContext,
+    embedding_matrix,
+    indices,
+    offsets,
+    scale_grad_by_freq,
+    mode,
+    sparse,
+    per_sample_weights,
+    include_last_offset,
+    padding_idx,
+):
+    return symbolic_helper._embedding_bag_helper(
+        g,
+        embedding_matrix,
+        indices,
+        offsets,
+        scale_grad_by_freq,
+        mode,
+        sparse,
+        per_sample_weights,
+        include_last_offset,
+        padding_idx,
+    )
+
+
+@_onnx_symbolic("aten::linalg_vector_norm")
+@symbolic_helper.parse_args("v", "f", "is", "b", "v")
+def linalg_vector_norm(
+    g: jit_utils.GraphContext,
+    self: torch._C.Value,
+    ord: float,
+    dim: Sequence[int] | None,
+    keepdim: bool,
+    dtype: torch._C.Value,
+):
+    return symbolic_helper._linalg_vector_norm_helper(g, self, ord, dim, keepdim, dtype)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset19.py

@@ -0,0 +1,31 @@
+"""This file exports ONNX ops for opset 19.
+
+Note [ONNX Operators that are added/updated in opset 19]
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-19-of-the-default-onnx-operator-set
+New operators:
+AveragePool
+Cast
+CastLike
+Constant
+DeformConv
+DequantizeLinear
+Equal
+Identity
+If
+Loop
+Pad
+QuantizeLinear
+Reshape
+Resize
+Scan
+Shape
+Size
+"""
+
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in symbolic_helper.py
+
+__all__: list[str] = []

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset20.py

@@ -0,0 +1,95 @@
+# mypy: allow-untyped-defs
+"""This file exports ONNX ops for opset 20.
+
+Note [ONNX Operators that are added/updated in opset 20]
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-20-of-the-default-onnx-operator-set
+New operators:
+    AffineGrid
+    ConstantOfShape
+    DFT
+    Gelu
+    GridSample
+    ImageDecoder
+    IsInf
+    IsNaN
+    ReduceMax
+    ReduceMin
+    RegexFullMatch
+    StringConcat
+    StringSplit
+"""
+
+import functools
+
+import torch.nn.functional as F
+from torch import _C
+from torch.onnx._internal.torchscript_exporter import (
+    jit_utils,
+    registration,
+    symbolic_helper,
+)
+
+
+# EDITING THIS FILE? READ THIS FIRST!
+# see Note [Edit Symbolic Files] in symbolic_helper.py
+
+__all__ = ["_grid_sampler", "_affine_grid_generator", "gelu"]
+
+
+def convert_grid_sample_mode(mode_s):
+    return (
+        "linear" if mode_s == "bilinear" else "cubic" if mode_s == "bicubic" else mode_s
+    )
+
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=20)
+
+
+@_onnx_symbolic("aten::grid_sampler")
+@symbolic_helper.parse_args("v", "v", "i", "i", "b")
+def _grid_sampler(
+    g: jit_utils.GraphContext,
+    input: _C.Value,
+    grid: _C.Value,
+    mode_enum: int,
+    padding_mode_enum: int,
+    align_corners: bool,
+):
+    mode_s = {v: k for k, v in F.GRID_SAMPLE_INTERPOLATION_MODES.items()}[mode_enum]  # type: ignore[call-arg, index]
+    # mode string changes at https://onnx.ai/onnx/operators/text_diff_GridSample_16_20.html
+    mode_s = convert_grid_sample_mode(mode_s)
+    padding_mode_s = {v: k for k, v in F.GRID_SAMPLE_PADDING_MODES.items()}[  # type: ignore[call-arg, index]
+        padding_mode_enum  # type: ignore[index]
+    ]
+    return g.op(
+        "GridSample",
+        input,
+        grid,
+        align_corners_i=int(align_corners),
+        mode_s=mode_s,
+        padding_mode_s=padding_mode_s,
+    )
+
+
+@_onnx_symbolic("aten::affine_grid_generator")
+@symbolic_helper.parse_args("v", "v", "b")
+def _affine_grid_generator(
+    g: jit_utils.GraphContext,
+    theta: _C.Value,
+    size: _C.Value,
+    align_corners: bool,
+):
+    return g.op(
+        "AffineGrid",
+        theta,
+        size,
+        align_corners_i=int(align_corners),
+    )
+
+
+@_onnx_symbolic("aten::gelu")
+@symbolic_helper.parse_args("v", "s")
+def gelu(g: jit_utils.GraphContext, self: _C.Value, approximate: str = "none"):
+    return g.op("Gelu", self, approximate_s=approximate)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset7.py

@@ -0,0 +1,73 @@
+# mypy: allow-untyped-defs
+"""
+Note [ONNX operators that are added/updated from opset 7 to opset 8]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+New operators:
+  Expand
+
+Updated operators:
+  Min, Max, Sum, Mean: supports multidirectional broadcasting.
+  MaxPool: added optional indices output.
+  Scan
+"""
+
+import functools
+import warnings
+
+from torch.onnx._internal.torchscript_exporter import (
+    jit_utils,
+    registration,
+    symbolic_helper,
+    symbolic_opset9 as opset9,
+)
+
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=7)
+
+block_listed_operators = (
+    "scan",
+    "expand",
+    "expand_as",
+    "meshgrid",
+    "adaptive_max_pool1d",
+    "adaptive_max_pool2d",
+    "adaptive_max_pool3d",
+    "max_pool1d_with_indices",
+    "max_pool2d_with_indices",
+    "max_pool3d_with_indices",
+)
+
+
+# NOTE: max, min, sum, mean: broadcasting is not supported in opset 7.
+# torch.max (same for torch.min) actually has two interfaces smashed together:
+# torch.max(x, dim, keepdim) and torch.max(x, y)
+@_onnx_symbolic("aten::max")
+def max(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None):
+    # torch.max(input, other)
+    if keepdim is None and dim_or_y is not None:
+        warnings.warn(
+            "Multidirectional broadcasting is not supported in opset 7. "
+            "This might cause the onnx model to be incorrect, if inputs to max operators "
+            "have different shapes",
+            stacklevel=2,
+        )
+    return opset9.max(g, self, dim_or_y, keepdim)
+
+
+@_onnx_symbolic("aten::min")
+def min(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None):
+    # torch.min(input, other)
+    if keepdim is None and dim_or_y is not None:
+        warnings.warn(
+            "Multidirectional broadcasting is not supported in opset 7. "
+            "This might cause the onnx model to be incorrect, if inputs to min operators "
+            "have different shapes",
+            stacklevel=2,
+        )
+    return opset9.min(g, self, dim_or_y, keepdim)
+
+
+for block_listed_op in block_listed_operators:
+    _onnx_symbolic(f"aten::{block_listed_op}")(
+        symbolic_helper._block_list_in_opset(block_listed_op)
+    )

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset8.py

@@ -0,0 +1,470 @@
+# mypy: allow-untyped-defs
+"""
+Note [ONNX operators that are added/updated from opset 8 to opset 9]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+New operators:
+    Compress
+    ConstantOfShape
+    EyeLike
+    MaxUnpool
+    OneHot
+    Sinh
+    Cosh
+    Asinh
+    Acosh
+    Atanh
+    Shrink
+    IsNaN
+    Sign
+    Erf
+    Scatter
+    Where
+    NonZero
+    TfIdfVectorizer
+    MeanVarianceNormalization
+
+Updated operators:
+    BatchNormalization: removed spatial attribute.
+    Greater, Less, Constant, MatMul, PRelu, Gemm, Flatten: more data types{integers} supported.
+    Cast: more data types{string} supported.
+    Upsample: moved scales from attribute to input.
+    Scan
+"""
+
+import functools
+import warnings
+
+import torch
+from torch._C import _onnx as _C_onnx
+from torch.onnx import errors
+from torch.onnx._internal.torchscript_exporter import (
+    _type_utils,
+    jit_utils,
+    registration,
+    symbolic_helper,
+    symbolic_opset9 as opset9,
+)
+
+
+_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=8)
+
+block_listed_operators = (
+    "nonzero",
+    "where",
+    "scatter",
+    "scatter_add",
+    "erf",
+    "sign",
+    "isnan",
+    "gather",
+    "arange",
+    "masked_fill",
+    "index_fill",
+    "index_copy",
+    "repeat_interleave",
+    "any",
+    "all",
+)
+
+for block_listed_op in block_listed_operators:
+    _onnx_symbolic(f"aten::{block_listed_op}")(
+        symbolic_helper._block_list_in_opset(block_listed_op)
+    )
+
+
+@_onnx_symbolic(
+    "aten::upsample_nearest1d",
+    decorate=[symbolic_helper._apply_params("upsample_nearest1d", 3, "nearest")],
+)
+@_onnx_symbolic(
+    "aten::upsample_nearest2d",
+    decorate=[symbolic_helper._apply_params("upsample_nearest2d", 4, "nearest")],
+)
+@_onnx_symbolic(
+    "aten::upsample_nearest3d",
+    decorate=[symbolic_helper._apply_params("upsample_nearest3d", 5, "nearest")],
+)
+@_onnx_symbolic(
+    "aten::upsample_linear1d",
+    decorate=[symbolic_helper._apply_params("upsample_linear1d", 3, "linear")],
+)
+@_onnx_symbolic(
+    "aten::upsample_bilinear2d",
+    decorate=[symbolic_helper._apply_params("upsample_bilinear2d", 4, "linear")],
+)
+@_onnx_symbolic(
+    "aten::upsample_trilinear3d",
+    decorate=[symbolic_helper._apply_params("upsample_trilinear3d", 5, "linear")],
+)
+def _interpolate(name, dim, interpolate_mode):
+    def symbolic_fn(g, input, output_size, *args):
+        scales, align_corners = symbolic_helper._get_interpolate_attributes(
+            g, interpolate_mode, args
+        )
+        symbolic_helper._interpolate_warning(interpolate_mode)
+        align_corners = symbolic_helper._maybe_get_scalar(align_corners)
+        if align_corners:
+            return symbolic_helper._unimplemented(name, "align_corners == True", input)
+        output_size = symbolic_helper._maybe_get_const(output_size, "is")
+        if symbolic_helper._is_value(output_size):
+            return symbolic_helper._unimplemented(
+                name, "torch._C.Value (output_size) indexing"
+            )
+        if scales is None:
+            scales = [
+                1.0
+                if i < 2
+                else float(output_size[-(dim - i)])
+                / float(input.type().sizes()[-(dim - i)])
+                for i in range(dim)
+            ]
+        return g.op("Upsample", input, mode_s=interpolate_mode, scales_f=scales)
+
+    return symbolic_fn
+
+
+@_onnx_symbolic("aten::__interpolate")
+def __interpolate(
+    g: jit_utils.GraphContext,
+    input,
+    size,
+    scale_factor,
+    mode,
+    align_corners,
+    recompute_scale_factor,
+    antialias,
+):
+    align_corners = symbolic_helper._maybe_get_const(align_corners, "b")
+    if not symbolic_helper._is_none(align_corners) and align_corners:
+        return symbolic_helper._unimplemented("interpolate", "align_corners == True")
+
+    if not symbolic_helper._is_none(scale_factor) and symbolic_helper._is_value(
+        scale_factor
+    ):
+        return symbolic_helper._unimplemented(
+            "interpolate", "dynamic scales in opset 8"
+        )
+
+    if not symbolic_helper._is_none(size) and symbolic_helper._is_value(size):
+        return symbolic_helper._unimplemented("interpolate", "dynamic size in opset 8")
+
+    scales, mode = symbolic_helper._interpolate_get_scales_and_mode(
+        g, input, size, scale_factor, mode, align_corners
+    )
+    return g.op("Upsample", input, mode_s=mode, scales_f=scales)
+
+
+# NOTE: We should create a wrapper for this kind of operation, after resolving the shape/type propagation
+#       issue for "cast" operators. Some symbolic functions depend on shape information of input tensor, which
+#       is lost after casting.
+def _try_cast_integer_to_float(g: jit_utils.GraphContext, *args):
+    floating_scalar_types = {
+        _type_utils.JitScalarType.HALF,
+        _type_utils.JitScalarType.FLOAT,
+        _type_utils.JitScalarType.DOUBLE,
+    }
+    old_type = None
+    # Cast the input tensor to Float if its scalarType is known and is not floating number.
+    # If casting is performed, return the old scalarType, otherwise return None.
+    arg0_type = _type_utils.JitScalarType.from_value(
+        args[0], _type_utils.JitScalarType.UNDEFINED
+    )
+    if arg0_type != _type_utils.JitScalarType.UNDEFINED:
+        old_type = arg0_type
+        if old_type not in floating_scalar_types:
+            old_type = old_type.scalar_name()  # type: ignore[assignment]
+            args = tuple(
+                g.op("Cast", arg, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+                for arg in args
+            )
+        else:
+            return (None,) + args
+    else:
+        warnings.warn(
+            "Only floating datatype is supported for these operators: "
+            "{Greater, Less, MatMul, PRelu, Gemm, Flatten}. This might cause "
+            "the onnx model to be incorrect, if inputs have integer datatypes.",
+            stacklevel=2,
+        )
+    return (old_type,) + args
+
+
+def _cast_to_type(g: jit_utils.GraphContext, input, to_type):
+    if to_type is None:
+        return input
+    return g.op("Cast", input, to_i=symbolic_helper.cast_pytorch_to_onnx[to_type])
+
+
+def _comparison_operator(g: jit_utils.GraphContext, input, other, op_name):
+    other = symbolic_helper._maybe_get_scalar(other)
+    other = symbolic_helper._if_scalar_type_as(other, input)
+    _, input, other = _try_cast_integer_to_float(g, input, other)
+    return g.op(op_name, input, other)
+
+
+# NOTE: For symbolics {gt, lt, bmm, matmul, prelu, mm, addmm, view, flatten},
+#       integer input type not supported in opset8. Cast to float if possible.
+@_onnx_symbolic("aten::gt")
+def gt(g: jit_utils.GraphContext, input, other):
+    return _comparison_operator(g, input, other, "Greater")
+
+
+@_onnx_symbolic("aten::lt")
+def lt(g: jit_utils.GraphContext, input, other):
+    return _comparison_operator(g, input, other, "Less")
+
+
+@_onnx_symbolic("aten::bmm")
+def bmm(g: jit_utils.GraphContext, self, other):
+    if symbolic_helper._try_get_scalar_type(self):
+        old_type, self, other = _try_cast_integer_to_float(g, self, other)
+        return _cast_to_type(g, g.op("MatMul", self, other), old_type)
+    else:
+        return g.op("MatMul", self, other)
+
+
+@_onnx_symbolic("aten::matmul")
+def matmul(g: jit_utils.GraphContext, self, other):
+    return bmm(g, self, other)
+
+
+@_onnx_symbolic("aten::prelu")
+def prelu(g: jit_utils.GraphContext, self, weight):
+    self_rank = symbolic_helper._get_tensor_rank(self)
+    weight_sizes = symbolic_helper._get_tensor_sizes(weight)
+    if self_rank is not None and self_rank > 2:
+        weight = g.op("Unsqueeze", weight, axes_i=list(range(1, self_rank - 1)))
+    elif self_rank == 0 and weight_sizes == [1]:
+        # self and weight are both scalar but weight has rank == 1, squeeze weight.
+        weight = symbolic_helper._squeeze_helper(g, weight, [0])
+    if symbolic_helper._try_get_scalar_type(self):
+        old_type, self, weight = _try_cast_integer_to_float(g, self, weight)
+        return _cast_to_type(g, g.op("PRelu", self, weight), old_type)
+    else:
+        return g.op("PRelu", self, weight)
+
+
+@_onnx_symbolic("aten::mm")
+def mm(g: jit_utils.GraphContext, self, other):
+    # Create a dummy C tensor. Only needed for API purposes, the value is
+    # since beta = 0
+    scalar_type = symbolic_helper._try_get_scalar_type(self, other)
+    if scalar_type is None:
+        raise errors.SymbolicValueError(
+            "mm can only operate on tensors with known types", self
+        )
+    zero_constant = g.op(
+        "Constant",
+        value_t=torch.tensor([0], dtype=scalar_type.dtype()),
+    )
+
+    if symbolic_helper._try_get_scalar_type(self):
+        old_type, self, other, zero_constant = _try_cast_integer_to_float(
+            g, self, other, zero_constant
+        )
+        return _cast_to_type(
+            g,
+            g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0),
+            old_type,
+        )
+    return g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0)
+
+
+@_onnx_symbolic("aten::addmm")
+@symbolic_helper.parse_args("v", "v", "v", "t", "t")
+def addmm(g: jit_utils.GraphContext, self, mat1, mat2, beta, alpha):
+    if symbolic_helper._try_get_scalar_type(self):
+        old_type, self, mat1, mat2 = _try_cast_integer_to_float(g, self, mat1, mat2)
+        return _cast_to_type(
+            g,
+            g.op(
+                "Gemm",
+                mat1,
+                mat2,
+                self,
+                beta_f=symbolic_helper._scalar(beta),
+                alpha_f=symbolic_helper._scalar(alpha),
+            ),
+            old_type,
+        )
+    else:
+        return g.op(
+            "Gemm",
+            mat1,
+            mat2,
+            self,
+            beta_f=symbolic_helper._scalar(beta),
+            alpha_f=symbolic_helper._scalar(alpha),
+        )
+
+
+@_onnx_symbolic("aten::flatten")
+def flatten(g: jit_utils.GraphContext, input, start_dim, end_dim):
+    start_dim_i = symbolic_helper._get_const(start_dim, "i", "start_dim")
+    end_dim_i = symbolic_helper._get_const(end_dim, "i", "end_dim")
+
+    dim = input.type().dim()
+    if end_dim_i < 0:
+        end_dim_i = dim + end_dim_i
+    # use ONNX's Flatten operator for cases where the output shape is 2D
+    if start_dim_i == 1 and end_dim_i == dim - 1:
+        if symbolic_helper._try_get_scalar_type(input):
+            old_type, input = _try_cast_integer_to_float(g, input)
+            return _cast_to_type(
+                g, g.op("Flatten", input, axis_i=start_dim_i), old_type
+            )
+        else:
+            return g.op("Flatten", input, axis_i=start_dim_i)
+    if start_dim_i == 0 and end_dim_i == dim - 2:
+        if symbolic_helper._try_get_scalar_type(input):
+            old_type, input = _try_cast_integer_to_float(g, input)
+            return _cast_to_type(
+                g, g.op("Flatten", input, axis_i=end_dim_i + 1), old_type
+            )
+        else:
+            return g.op("Flatten", input, axis_i=end_dim_i + 1)
+
+    return opset9.flatten(g, input, start_dim, end_dim)
+
+
+def _constant_fill(g: jit_utils.GraphContext, sizes, dtype: int, const_value):
+    if dtype is None:
+        scalar_type = _type_utils.JitScalarType.FLOAT
+    else:
+        scalar_type = _type_utils.JitScalarType(dtype)
+    if not scalar_type.dtype().is_floating_point:
+        result = g.op(
+            "ConstantFill",
+            sizes,
+            dtype_i=_type_utils.JitScalarType.FLOAT.onnx_type(),
+            input_as_shape_i=1,
+            value_f=const_value,
+        )
+        return g.op("Cast", result, to_i=scalar_type.onnx_type())
+    else:
+        return g.op(
+            "ConstantFill",
+            sizes,
+            dtype_i=scalar_type.onnx_type(),
+            input_as_shape_i=1,
+            value_f=const_value,
+        )
+
+
+@_onnx_symbolic("aten::empty")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v")
+def empty(
+    g: jit_utils.GraphContext,
+    sizes,
+    dtype,
+    layout,
+    device,
+    pin_memory=False,
+    memory_format=None,
+):
+    return zeros(g, sizes, dtype, layout, device, pin_memory)
+
+
+@_onnx_symbolic("aten::empty_like")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v")
+def empty_like(
+    g: jit_utils.GraphContext,
+    input,
+    dtype,
+    layout,
+    device,
+    pin_memory=False,
+    memory_format=None,
+):
+    return zeros_like(g, input, dtype, layout, device, pin_memory)
+
+
+@_onnx_symbolic("aten::zeros")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v")
+def zeros(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False):
+    # NOTE: no way to set device and layout in ONNX, so we ignore it
+    return _constant_fill(g, sizes, dtype, 0)
+
+
+@_onnx_symbolic("aten::zeros_like")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v")
+def zeros_like(
+    g: jit_utils.GraphContext,
+    input,
+    dtype,
+    layout,
+    device,
+    pin_memory=False,
+    memory_format=None,
+):
+    shape = g.op("Shape", input)
+    return _constant_fill(g, shape, dtype, 0)
+
+
+@_onnx_symbolic("aten::ones")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v")
+def ones(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False):
+    return _constant_fill(g, sizes, dtype, 1)
+
+
+@_onnx_symbolic("aten::ones_like")
+@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v")
+def ones_like(
+    g: jit_utils.GraphContext,
+    input,
+    dtype,
+    layout,
+    device,
+    pin_memory=False,
+    memory_format=None,
+):
+    shape = g.op("Shape", input)
+    return _constant_fill(g, shape, dtype, 1)
+
+
+@_onnx_symbolic("aten::full")
+def full(
+    g: jit_utils.GraphContext, sizes, value, dtype, layout, device, pin_memory=False
+):
+    const_value = symbolic_helper._maybe_get_const(value, "t")
+    if symbolic_helper._is_value(const_value):
+        tmp = zeros(g, sizes, dtype, layout, device)
+        return opset9.add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1)))
+    else:
+        dtype = symbolic_helper._get_const(dtype, "i", "dtype")
+        return _constant_fill(g, sizes, dtype, const_value)
+
+
+@_onnx_symbolic("aten::full_like")
+@symbolic_helper.parse_args("v", "f", "i", "v", "v", "v", "v")
+def full_like(
+    g: jit_utils.GraphContext,
+    input,
+    fill_value,
+    dtype,
+    layout,
+    device,
+    pin_memory=False,
+    memory_format=None,
+):
+    shape = g.op("Shape", input)
+    return _constant_fill(g, shape, dtype, fill_value)
+
+
+@_onnx_symbolic("aten::repeat")
+def repeat(g: jit_utils.GraphContext, self, repeats):
+    if not symbolic_helper._is_value(repeats):
+        repeats = g.op("Constant", value_t=torch.LongTensor(repeats))
+    if symbolic_helper._is_packed_list(repeats):
+        repeat_size_len = len(symbolic_helper._unpack_list(repeats))
+    else:
+        const_repeats = symbolic_helper._maybe_get_const(repeats, "is")
+        repeat_size_len = len(const_repeats)
+    if self.isCompleteTensor():
+        sizes = self.type().sizes()
+        diff_dims = repeat_size_len - len(sizes)
+        if diff_dims > 0:
+            self = opset9.view(
+                g, self, g.op("Constant", value_t=torch.tensor([1] * diff_dims + sizes))
+            )
+    return g.op("Tile", self, repeats)

kronos-venv/lib/python3.12/site-packages/torch/onnx/_internal/torchscript_exporter/verification.py

@@ -0,0 +1,516 @@
+# mypy: allow-untyped-defs
+"""The ONNX verification module provides a set of tools to verify the correctness of ONNX models."""
+
+from __future__ import annotations
+
+
+__all__ = [
+    "OnnxBackend",
+    "VerificationOptions",
+    "verify",
+]
+
+import contextlib
+import copy
+import dataclasses
+import enum
+import io
+import os
+import tempfile
+import warnings
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+import torch
+import torch._C._onnx as _C_onnx
+from torch.onnx._internal.torchscript_exporter import utils
+from torch.types import Number
+
+
+# Everything below are deprecated ##############################################
+
+_ORT_PROVIDERS = ("CPUExecutionProvider",)
+
+_NumericType = Number | torch.Tensor | np.ndarray
+_ModelType = torch.nn.Module | torch.jit.ScriptModule
+_InputArgsType = torch.Tensor | tuple[Any, ...]
+_InputKwargsType = Mapping[str, Any]
+_OutputsType = Sequence[_NumericType] | Sequence
+
+
+class OnnxBackend(enum.Enum):
+    """Enum class for ONNX backend used for export verification.
+
+    .. deprecated:: 2.7
+        Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned
+        ``ONNXProgram`` to test the ONNX model.
+    """
+
+    REFERENCE = "ONNXReferenceEvaluator"
+    ONNX_RUNTIME_CPU = "CPUExecutionProvider"
+    ONNX_RUNTIME_CUDA = "CUDAExecutionProvider"
+
+
+@dataclasses.dataclass
+class VerificationOptions:
+    """Options for ONNX export verification.
+
+    .. deprecated:: 2.7
+        Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned
+        ``ONNXProgram`` to test the ONNX model.
+
+    Attributes:
+        flatten: If True, unpack nested list/tuple/dict inputs into a flattened list of
+            Tensors for ONNX. Set this to False if nested structures are to be preserved
+            for ONNX, which is usually the case with exporting ScriptModules. Default True.
+        ignore_none: Whether to ignore None type in torch output, which is usually the
+            case with tracing. Set this to False, if torch output should keep None type,
+            which is usually the case with exporting ScriptModules. Default to True.
+        check_shape: Whether to check the shapes between PyTorch and ONNX Runtime outputs
+            are exactly the same. Set this to False to allow output shape broadcasting.
+            Default to True.
+        check_dtype: Whether to check the dtypes between PyTorch and ONNX Runtime outputs
+            are consistent. Default to True.
+        backend: ONNX backend for verification. Default to OnnxBackend.ONNX_RUNTIME_CPU.
+        rtol: relative tolerance in comparison between ONNX and PyTorch outputs.
+        atol: absolute tolerance in comparison between ONNX and PyTorch outputs.
+        remained_onnx_input_idx: If provided, only the specified inputs will be passed
+            to the ONNX model. Supply a list when there are unused inputs in the model.
+            Since unused inputs will be removed in the exported ONNX model, supplying
+            all inputs will cause an error on unexpected inputs. This parameter tells
+            the verifier which inputs to pass into the ONNX model.
+        acceptable_error_percentage: acceptable percentage of element mismatches in comparison.
+            It should be a float of value between 0.0 and 1.0.
+    """
+
+    flatten: bool = True
+    ignore_none: bool = True
+    check_shape: bool = True
+    check_dtype: bool = True
+    backend: OnnxBackend = OnnxBackend.ONNX_RUNTIME_CPU
+    rtol: float = 1e-3
+    atol: float = 1e-7
+    remained_onnx_input_idx: Sequence[int] | None = None
+    acceptable_error_percentage: float | None = None
+
+
+def _flatten_tuples(elem):
+    flattened = []
+    for t in elem:
+        if isinstance(t, tuple):
+            flattened.extend(_flatten_tuples(t))
+        else:
+            flattened.append(t)
+    return flattened
+
+
+def _to_numpy(elem) -> list | npt.NDArray:
+    if isinstance(elem, torch.Tensor):
+        if elem.requires_grad:
+            return elem.detach().cpu().numpy()
+        else:
+            return elem.cpu().numpy()
+    elif isinstance(elem, (list, tuple)):
+        return [_to_numpy(inp) for inp in elem]
+    elif isinstance(elem, (bool, int, float)):
+        return np.array(elem)
+    elif isinstance(elem, dict):
+        flattened = []
+        for k in elem:
+            flattened.extend([_to_numpy(k), _to_numpy(elem[k])])
+        return flattened
+    return elem
+
+
+def _inline_flatten_list(inputs, res_list) -> list:
+    for i in inputs:
+        res_list.append(i) if not isinstance(
+            i, (list, tuple)
+        ) else _inline_flatten_list(i, res_list)
+    return res_list
+
+
+def _unpack_to_numpy(values, cast_onnx_accepted=True) -> list:
+    value_unpacked = []
+    for value in values:
+        value_unpacked.extend(
+            utils.unpack_quantized_tensor(value, cast_onnx_accepted=cast_onnx_accepted)
+        )
+    return [_to_numpy(v) for v in value_unpacked]
+
+
+def _run_onnx(onnx_session, inputs) -> _OutputsType:
+    kw_inputs = {}
+    if inputs and isinstance(inputs[-1], dict):
+        kw_inputs = inputs[-1]
+        inputs = inputs[:-1]
+    inputs = _unpack_to_numpy(_flatten_tuples(inputs))
+    ort_inputs = {}
+    for input_name, input in kw_inputs.items():
+        ort_inputs[input_name] = _to_numpy(input)
+    inputs = _to_numpy(inputs)
+    if hasattr(onnx_session, "get_inputs"):
+        # onnxruntime.InferenceSession
+        input_names = [i.name for i in onnx_session.get_inputs()]
+    elif hasattr(onnx_session, "input_names"):
+        # onnx.reference.ReferenceEvaluator
+        input_names = onnx_session.input_names
+    else:
+        raise ValueError(f"Unknown ONNX backend type: {type(onnx_session)}.")
+
+    for i, input in enumerate(inputs):
+        if i == len(input_names) or input_names[i] in ort_inputs:
+            raise ValueError(
+                f"got too many positional inputs. inputs: {inputs}. kw_inputs: {kw_inputs}. "
+                f"input names: {input_names}."
+            )
+        ort_inputs[input_names[i]] = input
+    onnx_outs = onnx_session.run(None, ort_inputs)
+    return onnx_outs
+
+
+def _ort_session(
+    model: str | io.BytesIO, ort_providers: Sequence[str] = _ORT_PROVIDERS
+):
+    try:
+        import onnxruntime  # type: ignore[import]
+    except ImportError as e:
+        raise ImportError("onnxruntime is required for export verification.") from e
+
+    if ort_providers is None:
+        ort_providers = _ORT_PROVIDERS
+
+    session_options = onnxruntime.SessionOptions()
+    # suppress ort warnings.
+    # 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2.
+    session_options.log_severity_level = 3
+    ort_session = onnxruntime.InferenceSession(
+        model if isinstance(model, str) else model.getvalue(),
+        session_options,
+        providers=ort_providers,
+    )
+    return ort_session
+
+
+def _onnx_backend_session(model: str | io.BytesIO, backend: OnnxBackend):
+    if backend == OnnxBackend.REFERENCE:
+        raise NotImplementedError
+    elif backend in {OnnxBackend.ONNX_RUNTIME_CPU, OnnxBackend.ONNX_RUNTIME_CUDA}:
+        onnx_session = _ort_session(model, (backend.value,))
+    else:
+        raise ValueError(f"Unsupported backend: {backend}")
+    return onnx_session
+
+
+def _compare_onnx_pytorch_outputs_in_np(
+    onnx_outs: _OutputsType,
+    pt_outs: _OutputsType,
+    options: VerificationOptions,
+) -> None:
+    if len(onnx_outs) != len(pt_outs):
+        raise AssertionError(
+            f"Number of outputs differ ONNX runtime: ({len(onnx_outs)}) PyTorch: ({len(pt_outs)})"
+        )
+    acceptable_error_percentage = options.acceptable_error_percentage
+    if acceptable_error_percentage and (
+        acceptable_error_percentage > 1.0 or acceptable_error_percentage < 0.0
+    ):
+        raise ValueError(
+            "If set, acceptable_error_percentage should be between 0.0 and 1.0"
+        )
+
+    for ort_out, pt_out in zip(onnx_outs, pt_outs):
+        try:
+            # TODO: Remove `check_shape` option once every shape inconsistent issue is addressed.
+            if not options.check_shape:
+                # Allow different but broadcastable output shapes.
+                ort_out, pt_out = np.broadcast_arrays(ort_out, pt_out)
+            torch.testing.assert_close(
+                ort_out,
+                pt_out,
+                rtol=options.rtol,
+                atol=options.atol,
+                check_dtype=options.check_dtype,
+                equal_nan=True,
+            )
+        except AssertionError as e:
+            if acceptable_error_percentage:
+                error_percentage = 1 - np.sum(
+                    np.isclose(ort_out, pt_out, rtol=options.rtol, atol=options.atol)
+                ) / np.prod(ort_out.shape)  # pyrefly: ignore [missing-attribute]
+                if error_percentage <= acceptable_error_percentage:
+                    warnings.warn(
+                        f"Suppressed AssertionError:\n{e}.\n"
+                        f"Error percentage {error_percentage} "
+                        f"within acceptable range {acceptable_error_percentage}.",
+                        stacklevel=2,
+                    )
+                    continue
+            # pyrefly: ignore [missing-attribute]
+            if ort_out.dtype == np.uint8 or ort_out.dtype == np.int8:
+                warnings.warn("ONNX output is quantized", stacklevel=2)
+            # pyrefly: ignore [missing-attribute]
+            if pt_out.dtype == np.uint8 or pt_out.dtype == np.int8:
+                warnings.warn("PyTorch output is quantized", stacklevel=2)
+            raise
+
+
+def _compare_onnx_pytorch_outputs(
+    onnx_outs: _OutputsType,
+    pt_outs: Any,
+    options: VerificationOptions,
+) -> None:
+    """
+    Compare ONNX and PyTorch outputs.
+
+    Args:
+        onnx_outs: outputs from ONNX backend.
+        pt_outs: outputs from PyTorch.
+        options: options for verification.
+
+    Raises:
+        AssertionError: if outputs from ONNX model and PyTorch model are not
+            equal up to specified precision.
+        ValueError: if arguments provided are invalid.
+    """
+    if options.ignore_none:
+        # torch.jit._flatten filters None type
+        pt_outs, _ = torch.jit._flatten(pt_outs)
+    else:
+        pt_outs = _inline_flatten_list([pt_outs], [])
+    pt_outs_np = _unpack_to_numpy(pt_outs, cast_onnx_accepted=False)
+    onnx_outs = _inline_flatten_list(onnx_outs, [])
+    _compare_onnx_pytorch_outputs_in_np(onnx_outs, pt_outs_np, options)
+
+
+def _prepare_input_for_pytorch(args, kwargs):
+    """Prepare input for PyTorch model execution.
+
+    Any future changes/formatting to the input before dispatching to the PyTorch
+    model should be made in this function.
+
+    Args:
+        args: positional arguments for PyTorch model forward method.
+        kwargs: keyword arguments for PyTorch model forward method.
+
+    Returns:
+        args: positional arguments for PyTorch model forward method.
+        kwargs: keyword arguments for PyTorch model forward method.
+    """
+    if isinstance(args, (torch.Tensor, dict)):
+        args = (args,)
+    # In-place operators will update input tensor data as well.
+    # Thus inputs are replicated before every forward call.
+    args = copy.deepcopy(args)
+    if kwargs:
+        kwargs = copy.deepcopy(kwargs)
+    else:
+        kwargs = {}
+    return args, kwargs
+
+
+def _prepare_input_for_export(args, kwargs):
+    """Prepare input for ONNX model export.
+
+    Any future changes/formatting to the input before dispatching to the
+    :func:`torch.onnx.export` api should be made in this function.
+
+    Args:
+        args: positional arguments for PyTorch model forward method.
+        kwargs: keyword arguments for PyTorch model forward method.
+
+    Returns:
+        onnx_inputs: positional arguments for ONNX model export, as `args` in
+            :func:`torch.onnx.export`.
+    """
+    args, kwargs = _prepare_input_for_pytorch(args, kwargs)
+    if not kwargs and len(args) > 0 and isinstance(args[-1], dict):
+        onnx_inputs = args + ({},)
+    elif kwargs:
+        onnx_inputs = args + (kwargs,)
+    else:
+        onnx_inputs = args
+    return onnx_inputs
+
+
+def _prepare_input_for_onnx(
+    args, kwargs, remained_onnx_input_idx: Sequence[int] | None, flatten: bool
+):
+    """Prepare input for ONNX model execution in ONNX backend.
+
+    Any future changes/formatting to the input before dispatching to the ONNX backend
+    run should be made in this function.
+
+    Args:
+        args: positional arguments for PyTorch model forward method.
+        kwargs: keyword arguments for PyTorch model forward method.
+        remained_onnx_input_idx: indices of inputs to be used for ONNX model execution.
+        flatten: whether to flatten the input before dispatching to the ONNX model execution.
+
+    Returns:
+        onnx_inputs: positional arguments for ONNX model execution in ONNX backend.
+    """
+    onnx_inputs = _prepare_input_for_export(args, kwargs)
+    if flatten:
+        onnx_inputs, _ = torch.jit._flatten(onnx_inputs)
+    elif onnx_inputs and onnx_inputs[-1] == {}:
+        # Handle empty kwargs (normally removed by flatten).
+        onnx_inputs = onnx_inputs[:-1]
+    if remained_onnx_input_idx is not None:
+        return [onnx_inputs[i] for i in remained_onnx_input_idx]
+    else:
+        return onnx_inputs
+
+
+def _try_clone_model(model):
+    """Used for preserving original model in case forward mutates model states."""
+    try:
+        return copy.deepcopy(model)
+    except Exception:
+        warnings.warn(
+            "Failed to clone model. Model state might be mutated during verification.",
+            stacklevel=2,
+        )
+        return model
+
+
+def _compare_onnx_pytorch_model(
+    pt_model: _ModelType,
+    onnx_model_f: str | io.BytesIO,
+    input_args: _InputArgsType,
+    input_kwargs: _InputKwargsType | None,
+    additional_test_inputs: Sequence[_InputArgsType] | None,
+    options: VerificationOptions,
+) -> None:
+    """Compare outputs from ONNX model runs with outputs from PyTorch model runs.
+
+    Args:
+        pt_model: PyTorch model.
+        onnx_model_f: ONNX model file path or file-like object.
+        input_args: positional arguments for PyTorch model forward method.
+        input_kwargs: keyword arguments for PyTorch model forward method.
+        additional_test_inputs: additional positional arguments for PyTorch model
+            forward method.
+        options: options for verification.
+
+    Raises:
+        AssertionError: if outputs from ONNX model and PyTorch model are not
+            equal up to specified precision.
+    """
+    onnx_session = _onnx_backend_session(onnx_model_f, options.backend)
+
+    def compare_onnx_pytorch_model_with_input(input_args, input_kwargs) -> None:
+        pt_args, pt_kwargs = _prepare_input_for_pytorch(input_args, input_kwargs)
+        # TODO: remove this and treat mutating model separately. See #77679
+        pt_model_copy = _try_clone_model(pt_model)
+        pt_outs = pt_model_copy(*pt_args, **pt_kwargs)
+
+        onnx_inputs = _prepare_input_for_onnx(
+            input_args, input_kwargs, options.remained_onnx_input_idx, options.flatten
+        )
+
+        onnx_outs = _run_onnx(onnx_session, onnx_inputs)
+
+        _compare_onnx_pytorch_outputs(
+            onnx_outs=onnx_outs,
+            pt_outs=pt_outs,
+            options=options,
+        )
+
+    compare_onnx_pytorch_model_with_input(input_args, input_kwargs)
+
+    if additional_test_inputs:
+        for test_input_args in additional_test_inputs:
+            compare_onnx_pytorch_model_with_input(test_input_args, {})
+
+
+def verify(
+    model: _ModelType,
+    input_args: _InputArgsType,
+    input_kwargs: _InputKwargsType | None = None,
+    do_constant_folding: bool = True,
+    dynamic_axes: Mapping[str, Mapping[int, str] | Mapping[str, Sequence[int]]]
+    | None = None,
+    input_names: Sequence[str] | None = None,
+    output_names: Sequence[str] | None = None,
+    training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL,
+    opset_version: int | None = None,
+    keep_initializers_as_inputs: bool = True,
+    verbose: bool = False,
+    fixed_batch_size: bool = False,
+    use_external_data: bool = False,
+    additional_test_inputs: Sequence[_InputArgsType] | None = None,
+    options: VerificationOptions | None = None,
+) -> None:
+    """Verify model export to ONNX against original PyTorch model.
+
+    .. deprecated:: 2.7
+        Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned
+        ``ONNXProgram`` to test the ONNX model.
+
+    Args:
+        model: See :func:`torch.onnx.export`.
+        input_args: See :func:`torch.onnx.export`.
+        input_kwargs: See :func:`torch.onnx.export`.
+        do_constant_folding: See :func:`torch.onnx.export`.
+        dynamic_axes: See :func:`torch.onnx.export`.
+        input_names: See :func:`torch.onnx.export`.
+        output_names: See :func:`torch.onnx.export`.
+        training: See :func:`torch.onnx.export`.
+        opset_version: See :func:`torch.onnx.export`.
+        keep_initializers_as_inputs: See :func:`torch.onnx.export`.
+        verbose: See :func:`torch.onnx.export`.
+        fixed_batch_size: Legacy argument, used only by rnn test cases.
+        use_external_data: Explicitly specify whether to export the model with external data.
+        additional_test_inputs: List of tuples. Each tuple is a group of
+            input arguments to test. Currently only ``*args`` are supported.
+        options: A VerificationOptions object that controls the verification behavior.
+
+    Raises:
+        AssertionError: if outputs from ONNX model and PyTorch model are not
+            equal up to specified precision.
+        ValueError: if arguments provided are invalid.
+    """
+    if options is None:
+        options = VerificationOptions()
+
+    if training == torch.onnx.TrainingMode.TRAINING:
+        model.train()
+    elif training == torch.onnx.TrainingMode.EVAL:
+        model.eval()
+    with torch.no_grad(), contextlib.ExitStack() as stack:
+        model_f: str | io.BytesIO = io.BytesIO()
+        if use_external_data:
+            tmpdir_path = stack.enter_context(tempfile.TemporaryDirectory())
+            model_f = os.path.join(tmpdir_path, "model.onnx")
+
+        inputs_for_export = _prepare_input_for_export(input_args, input_kwargs)
+
+        # TODO(#77679): remove this and treat mutating model separately.
+        model_copy = _try_clone_model(model)
+        utils._export(
+            model,
+            inputs_for_export,
+            model_f,
+            opset_version=opset_version,
+            do_constant_folding=do_constant_folding,
+            keep_initializers_as_inputs=keep_initializers_as_inputs,
+            dynamic_axes=dynamic_axes,
+            input_names=input_names,
+            output_names=output_names,
+            fixed_batch_size=fixed_batch_size,
+            training=training,
+            verbose=verbose,
+        )
+
+        _compare_onnx_pytorch_model(
+            pt_model=model_copy,
+            onnx_model_f=model_f,
+            input_args=input_args,
+            input_kwargs=input_kwargs,
+            additional_test_inputs=additional_test_inputs,
+            options=options,
+        )