_constant_folding.py

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# NOTE: This will eventually replace the existing constant_folding.py and evaluator.py files.
from __future__ import annotations

__all__ = [
    "basic_constant_propagation",
    "fold_constants",
    "FoldConstantsPass",
    "FOLDED_FROM_KEY",
]

import dataclasses
import logging
import math
import typing
from typing import Any, Callable, Iterable, Sequence, Union

import numpy as np
import onnx
import onnx.reference.ops
import onnx_ir as ir

import onnxscript.utils.utils as utils
from onnxscript.ir import _tape

DEFAULT_CONSTANT_FOLD_BLACKLIST = [
    # ConstantOfShape is preserved to avoid increasing model size unnecessarily
    "ConstantOfShape",
    # Quantize/DequantizeLinear are preserved to keep the quantization info
    "DequantizeLinear",
    "DynamicQuantizeLinear",
    "QuantizeLinear",
]

DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT = 8192

DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT = 512 * 512

# Key used to store the metadata
FOLDED_FROM_KEY = "pkg.onnxscript.optimizer.folded_from"


_NON_DETERMINISTIC_OPS = frozenset(
    {
        "RandomUniform",
        "RandomNormal",
        "RandomUniformLike",
        "RandomNormalLike",
        "Multinomial",
    }
)

# A list of ops to always fold regardless of their input size limits, as long as
# they are the single consumer of the large input tensors
_DEFAULT_ALWAYS_FOLD_OPS = frozenset(
    {
        ("", "Transpose"),
    }
)

logger = logging.getLogger(__name__)


def _is_control_flow_op(node: ir.Node) -> bool:
    graph_types = {ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS}
    return any(attr.type in graph_types for attr in node.attributes.values())


def _is_non_deterministic_op(node: ir.Node) -> bool:
    return node.op_type in _NON_DETERMINISTIC_OPS and utils.is_onnx_domain(node.domain)


def _is_onnx_op(node: ir.Node, op_type: str) -> bool:
    return node.op_type == op_type and utils.is_onnx_domain(node.domain)


# "Standard" evaluators are used to perform constant-folding.
# The API below works only for non-control-flow ops (ops without any graph-attributes).
# This currently used ONNX's reference implementation. But we could also
# use ORT's implementation if we want to.
def _process_constant_node(node: ir.Node) -> None:
    """Sets const_value of output value of a Constant op node."""
    if not _is_onnx_op(node, "Constant"):
        return
    if len(node.attributes) != 1:
        return
    attr_name, attr_value = next(iter(node.attributes.items()))
    if len(node.outputs) != 1:
        return
    ir_value = node.outputs[0]

    if attr_value is None or not isinstance(attr_value, ir.Attr):
        return

    # Even if this is an attribute, the value property might not be set, which
    # happens e.g. in case of attribute references, i.e., ref_attr_name is set
    if attr_value.value is None:
        # For now reject this to prevent TypeError from accessing Nones below
        return

    const_value: ir.TensorProtocol
    if attr_name in {"value_float", "value_floats"}:
        const_value = ir.Tensor(
            np.array(attr_value.value, dtype=np.float32), name=ir_value.name
        )
    elif attr_name in {"value_int", "value_ints"}:
        const_value = ir.Tensor(np.array(attr_value.value, dtype=np.int64), name=ir_value.name)
    elif attr_name in {"value_string", "value_strings"}:
        const_value = ir.StringTensor(
            np.array(attr_value.value, dtype=np.bytes_), name=ir_value.name
        )
    elif attr_name == "value":
        const_value = typing.cast(ir.TensorProtocol, attr_value.value)
    else:
        return

    ir_value.const_value = const_value
    ir_value.shape = const_value.shape  # type: ignore
    ir_value.dtype = const_value.dtype


def basic_constant_propagation(nodes: Iterable[ir.Node]) -> None:
    """Performs basic constant propagation for a sequence of nodes.
    Just marks the output values of Constant op nodes with their const_value.
    """
    for node in nodes:
        _process_constant_node(node)


class ReferenceEvaluator:
    def get_evaluator(self, domain: str, op: str, version: int) -> Callable | None:
        try:
            op_impl_class = onnx.reference.ops.load_op(domain, op, version)
            return op_impl_class.eval  # noqa: TRY300
        except Exception:
            return None

    def evaluate(self, domain: str, op: str, version: int, *args, **kwargs) -> Any:
        logger.debug("Evaluating %s::%s", domain, op)
        evaluator = self.get_evaluator(domain, op, version)
        if evaluator is None:
            return None
        try:
            return evaluator(*args, **kwargs)
        except Exception as e:
            logger.warning("Evaluation failed: %s", e)
            return None


_reference_evaluator = ReferenceEvaluator()


@dataclasses.dataclass
class Replacement:
    """A replacement for a node in the graph."""

    new_outputs: Sequence[ir.Value]
    new_nodes: Sequence[ir.Node]


# The optimizer tracks an optional symbolic value for each value in the model.
# The symbolic value attached to a value X can be:
# - another IR value Y (indicating that X is equal to Y)
# - a list of IR values [Y1, Y2, ...] (indicating that X is a sequence of values Y1, Y2, ...)
# - a Shape object (indicating that X is a shape value)
# A Shape object as a symbolic value indicates that the corresponding value is
# 1-D (or 0-D) tensor of INT64 values. The values in this object may be constants
# or symbolic dimension values (like "batch_size", "sequence_length", etc.).
# Currently, we assume that symbolic dimensions are also guaranteed to be non-negative.
# TODO: Add support for negative symbolic dimensions.

SymbolicValue = Union[ir.Value, list[ir.Value], ir.Shape]


class OptimizerState:
    def __init__(self):
        self._sym_value_map: dict[ir.Value, SymbolicValue] = {}
        self._initializer_inputs: list[set[ir.Value]] = []

    @property
    def symbolic_value_map(self) -> dict[ir.Value, SymbolicValue]:
        return self._sym_value_map

    def get_sym_value(self, value: ir.Value | None) -> SymbolicValue | None:
        if value is None:
            return None
        return self._sym_value_map.get(value)

    def set_sym_value(self, value: ir.Value, sym_value: SymbolicValue) -> None:
        self._sym_value_map[value] = sym_value

    def get_shape_value(self, value: ir.Value | None) -> ir.Shape | None:
        const_value = _get_numpy_value(value, ir.DataType.INT64, size_limit=10)
        if const_value is not None:
            if const_value.ndim == 1:
                return ir.Shape(const_value.tolist())
            return None
        sym_value = self.get_sym_value(value)
        if isinstance(sym_value, ir.Shape):
            return sym_value
        # TODO use shape of value if available
        return None


# The "partial evaluators" below are non-standard evaluators. They are used to perform
# partial evaluation and/or static program analysis (abstract interpretation).
# A partial-evaluator function takes a node, a RewriterContext, OptimizerState and returns
# a Replacement for the node or None (if no replacement is needed). It may also return just
# the ir.Value or ir.Values to replace the output values of the node, when the new nodes
# can be inferred from the RewriterContext used to build the new nodes.
RewriterContext = _tape.Builder
ReturnValue = Union[Replacement, Sequence[ir.Value], ir.Value, None]
PartialEvaluatorFunction = Callable[[ir.Node, RewriterContext, OptimizerState], ReturnValue]


@dataclasses.dataclass
class PartialEvaluator:
    """A class that represents a partial-evaluator for a particular op.

    It is applicable for a specific version range (min_version, max_version) of the op.
    The min_version and max_version can be None, indicating that there is no version
    constraint in that direction.
    """

    min_version: int | None
    max_version: int | None
    function: PartialEvaluatorFunction

    def valid_for(self, version: int) -> bool:
        """Returns True if this evaluator is applicable for the given version."""
        return (self.min_version is None or version >= self.min_version) and (
            self.max_version is None or version <= self.max_version
        )


class PartialEvaluatorRegistry:
    """A class that maintains a registry of evaluators for ops."""

    def __init__(self):
        self.op_evaluators: dict[tuple[str, str], list[PartialEvaluator]] = {}

    def lookup_evaluators(self, domain: str, opname: str, version: int):
        evaluator_list = self.op_evaluators.get((domain, opname), [])
        return [
            evaluator.function for evaluator in evaluator_list if evaluator.valid_for(version)
        ]

    def register(
        self, opname: str, domain: str = "", version=None
    ) -> Callable[[PartialEvaluatorFunction], PartialEvaluatorFunction]:
        if (domain, opname) in self.op_evaluators:
            evaluator_list = self.op_evaluators[(domain, opname)]
        else:
            evaluator_list = []
            self.op_evaluators[(domain, opname)] = evaluator_list
        if version is None:
            min_version = None
            max_version = None
        elif isinstance(version, int):
            min_version = version
            max_version = version
        elif isinstance(version, tuple):
            min_version, max_version = version

        def decorator(function: PartialEvaluatorFunction) -> PartialEvaluatorFunction:
            evaluator_list.append(PartialEvaluator(min_version, max_version, function))
            return function

        return decorator


registry: PartialEvaluatorRegistry = PartialEvaluatorRegistry()

register = registry.register


def _same_shape(shape1: ir.Shape, shape2: ir.Shape) -> bool:
    # Comparison of shapes as tuples works except if any dimension is None
    # (which represents an unknown dimension value). Thus, two shapes such
    # as (Batch, 1024) and (Batch, 1024) are considered equal, but (None, 1024)
    # and (None, 1024) are not considered equal.
    if any(isinstance(dim, ir.SymbolicDim) and dim.value is None for dim in shape1):
        return False
    return shape1.dims == shape2.dims


def _get_numpy_value(
    val: ir.Value | None, dtype: ir.DataType | None = None, size_limit: int | None = None
) -> np.ndarray | None:
    """Returns the numpy value of a constant value, if available.

    It returns None if the value is not a constant value, or if the value is not of
    the specified element dtype, or if the size of the value exceeds the specified
    size_limit.
    """
    if val is None:
        return None
    const_value = val.const_value
    if const_value is not None:
        if dtype is not None and const_value.dtype != dtype:
            return None
        if size_limit is not None and const_value.size > size_limit:
            return None
        try:
            # Turn the constant value into a numpy array representation with the
            # specifics of this conversion handled by the tensor type
            array = const_value.numpy()
            # Can/should not reinterpret strings via .view, resulting in
            #   "TypeError: Cannot change data-type for array of references."
            # There is also no reason to reinterpret strings, this is only
            # relevant for some arithmetic types
            if const_value.dtype != ir.DataType.STRING:
                # Reinterpret the array with `.view()` because some
                # implementations  of ir.TensorProtocol (e.g. PyTorch<=2.7) do
                # not use ml_dtypes for bfloat16 etc.
                array = array.view(const_value.dtype.numpy())
        except FileNotFoundError:
            # External data is not available.
            logger.warning(
                "External data for value '%s' is not available. "
                "This may lead to incorrect constant folding.",
                val.name,
            )
            return None
        assert isinstance(array, np.ndarray)
        return array
    return None


def _get_bool_value(val: ir.Value | None) -> bool | None:
    if val is None:
        return None
    value = _get_numpy_value(val)
    if value is None:
        return None
    if value.size == 1 and value.dtype == bool:
        return value.item(0)
    return None


def _get_input(node: ir.Node, index: int) -> ir.Value | None:
    if index < len(node.inputs):
        return node.inputs[index]
    return None


def _get_output(node: ir.Node, index: int) -> ir.Value | None:
    if index < len(node.outputs):
        return node.outputs[index]
    return None


def _get_input_element_type(node: ir.Node, index: int) -> int:
    input = _get_input(node, index)
    if input is not None and input.type is not None:
        return input.type.dtype.value
    return ir.DataType.UNDEFINED.value


def _get_int_attribute(node: ir.Node, name: str, default: int | None = None) -> int | None:
    if name in node.attributes:
        attr = node.attributes[name]
        if not isinstance(attr, ir.Attr):
            return None
        attr_val = attr.value
        if isinstance(attr_val, int):
            return attr_val
        # This is an invalid model: attribute has invalid/unexpected type.
        # For now, we just return None. We could raise an error too.
        return None
    return default


@register("Add")
def add(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Propagate symbolic dim values."""

    def get_dim_value(input_index):
        input = _get_input(node, input_index)
        if input is None:
            return None
        shape_value: ir.Shape | None = state.get_shape_value(input)
        if shape_value is None or len(shape_value) != 1:
            return None
        dim: int | ir.SymbolicDim = shape_value[0]
        return dim if isinstance(dim, int) else dim.value

    dim0 = get_dim_value(0)
    dim1 = get_dim_value(1)
    if dim0 is None or dim1 is None:
        return None
    if isinstance(dim0, int) and isinstance(dim1, int):
        result_dim_value: int | ir.SymbolicDim = dim0 + dim1
    else:
        result_dim_value = ir.SymbolicDim(f"{dim0}+{dim1}")
    output = _get_output(node, 0)
    if output is not None:
        state.set_sym_value(output, ir.Shape([result_dim_value]))


@register("Abs")
def abs(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace an Abs node by Identity when applicable.

    Currently, addresses Abs applied to symbolic shapes.
    """
    input = _get_input(node, 0)
    input_sym_value = state.get_shape_value(input)
    if input_sym_value is None:
        return None
    if any(isinstance(d, int) and d < 0 for d in input_sym_value):
        return None
    # Abs applied to a symbolic shape of the form [1, 1, SequenceLength].
    # We assume that SequenceLength is a non-negative integer.
    # The Abs op is redundant in this case.
    return op.Identity(input)


@register("Gather")
def gather(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace a Gather node by a constant when applicable.

    Currently, handles the case of Gathering from a shape tensor.
    """
    input = _get_input(node, 0)
    indices = _get_input(node, 1)
    if input is None or indices is None:
        return None
    input_sym_value = state.get_shape_value(input)
    if input_sym_value is None:
        return None
    axis = _get_int_attribute(node, "axis", None)
    if axis != 0:
        return None
    indices_numpy_value = _get_numpy_value(indices)
    if indices_numpy_value is None:
        return None
    if indices_numpy_value.ndim != 1:
        return None
    gathered = [input_sym_value[i] for i in indices_numpy_value]
    output = _get_output(node, 0)
    if output is not None:
        state.set_sym_value(output, ir.Shape(gathered))
    if all(isinstance(d, int) for d in gathered):
        return op.Constant(value_ints=ir.AttrInt64s("value_ints", gathered))
    return None


def _propagate_shape_value(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Propagates symbolic shape value of input 0 to output 0.

    Applies to ops like Reshape/Squeeze/Unsqueeze where the shape of the tensor may change
    but the values in the tensor remain the same.
    """
    input = _get_input(node, 0)
    input_shape_value = state.get_shape_value(input)
    output = _get_output(node, 0)
    if output is not None and input_shape_value is not None:
        state.set_sym_value(output, input_shape_value)
    return None


@register("Reshape")
def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace a Reshape node by Identity when applicable.
    Also propagate symbolic shape values.
    """
    input = _get_input(node, 0)
    shape = _get_input(node, 1)
    if input is None or shape is None:
        return None

    input_shape = input.shape
    shape_value = state.get_shape_value(shape)

    if shape_value is None or input_shape is None:
        return _propagate_shape_value(node, op, state)

    # No need to check for special values like -1, 0, etc. here
    if _same_shape(input_shape, shape_value):
        return op.Identity(input)
    return _propagate_shape_value(node, op, state)


@register("Squeeze")
def squeeze(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Propagate symbolic shape values."""
    return _propagate_shape_value(node, op, state)


@register("Cast")
def cast(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input = _get_input(node, 0)
    output = _get_output(node, 0)

    if input is None or output is None:
        return None

    input_dtype = _get_input_element_type(node, 0)
    output_dtype = _get_int_attribute(node, "to", None)
    if output_dtype is not None:
        if input_dtype == output_dtype:
            return op.Identity(input)
        output.type = ir.TensorType(ir.DataType(output_dtype))
    return None


@register("CastLike")
def cast_like(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input0 = node.inputs[0]
    source_element_type = _get_input_element_type(node, 0)
    target_element_type = _get_input_element_type(node, 1)

    if target_element_type == ir.DataType.UNDEFINED:
        return None
    if source_element_type == target_element_type:
        return op.Identity(input0)
    return op.Cast(input0, to=target_element_type)


@register("Shape")
def shape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input = node.inputs[0]
    if input is None:
        return None
    shape = input.shape
    if shape is None:
        return None
    start = _get_int_attribute(node, "start", 0)
    end = _get_int_attribute(node, "end", None)
    shape_slice = shape[start:end]
    output = _get_output(node, 0)
    if output is not None:
        state.set_sym_value(output, ir.Shape(shape_slice))
    if all(isinstance(d, int) for d in shape_slice):
        return op.Constant(value_ints=ir.AttrInt64s("value_ints", list(shape_slice)))
    return None


@register("Size")
def size(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input = _get_input(node, 0)
    if input is None:
        return None
    shape = input.shape
    if shape is None:
        return None
    size = 1
    for d in shape:
        if not isinstance(d, int):
            return None
        size *= d
    return op.Constant(value_int=size)


def _move_initializers_to_graph(src: ir.Graph, dst: ir.Graph) -> None:
    """Move all initializers from src graph to dst graph, ensuring name uniqueness.
    When an If branch is inlined into the main graph, the branch subgraph may
    hold initializers (e.g. a constant axes tensor for a Squeeze node) that were
    folded in a prior pass. Those initializers must be migrated to the main graph
    so that the inlined nodes can still reference them; failing to do so leaves the
    references dangling and produces an invalid model.
    """
    counter: dict[str, int] = {}
    for name in list(src.initializers):
        initializer = src.initializers.pop(name)
        # Ensure name uniqueness in the destination graph
        new_name = name
        while new_name in dst.initializers:
            counter[name] = counter.get(name, 0) + 1
            new_name = f"{name}_{counter[name]}"
        if new_name != name:
            initializer.name = new_name
        dst.register_initializer(initializer)


@register("If")
def if_op(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    cond_input = _get_input(node, 0)
    cond = _get_bool_value(cond_input)
    if cond is not None:
        # cond is a constant-value: inline the branch
        branch = "then_branch" if cond else "else_branch"
        graph_attr = node.attributes.get(branch)
        if graph_attr is None:
            return None
        if graph_attr.type != ir.AttributeType.GRAPH:
            return None
        assert isinstance(graph_attr, ir.Attr)
        graph = graph_attr.as_graph()
        # Copy the graph outputs and clear the graph outputs so that the values are free to move
        formal_outs = list(graph.outputs)
        graph.outputs.clear()
        actual_outs = node.outputs
        renamings = {
            formal.name: actual.name
            for formal, actual in zip(formal_outs, actual_outs)
            if actual is not None
        }
        # TODO: Extend renaming to intermediate values.
        def rename(name):
            return renamings.get(name, name)

        graph_nodes = list(graph)
        graph.remove(graph_nodes)
        for sub_node in graph_nodes:
            # TODO: handle renaming inside subgraphs in nodes
            for v in sub_node.outputs:
                v.name = rename(v.name)
            # Avoid name collision.
            sub_node.name = f"{node.name}_{sub_node.name}"

        # Move initializers from the subgraph to the main graph to avoid losing them.
        # When the If branch was processed in a prior constant-folding pass, any
        # constants inside the branch (e.g. the 'axes' tensor for a Squeeze node)
        # may have been folded into subgraph initializers. Without this step those
        # initializers would be orphaned once the branch nodes are inlined here.
        main_graph = node.graph
        if main_graph is not None:
            _move_initializers_to_graph(graph, main_graph)

        # TODO: we should handle initializers as well!
        return Replacement(formal_outs, graph_nodes)
    return None


@register("Identity")
def identity(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    del op
    input = node.inputs[0]
    output = node.outputs[0]
    if input is not None and output is not None:
        # NOTE: backward shape inference
        try:
            input.shape = _merge_shapes(input.shape, output.shape)
        except Exception as e:
            logger.warning(
                "[Constant folder] Cannot merge shapes on Identity node '%s' "
                "(folded from: %s) because of error: %s",
                node.name,
                input.meta.get(FOLDED_FROM_KEY, set()),
                e,
            )
        if input.type is None:
            input.type = output.type
        state.set_sym_value(output, input)
    return None


@register("SequenceConstruct")
def sequence_construct(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    del op
    output = node.outputs[0]
    if output is not None:
        state.set_sym_value(output, list(node.inputs))
    return None


@register("Concat")
def concat(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace a Concat node with a single input by Identity"""

    # Replace Concat(x) by Identity(x)
    inputs = node.inputs
    if len(inputs) == 1:
        return op.Identity(inputs[0])

    axis = _get_int_attribute(node, "axis", None)
    if axis is None:
        return None

    # Eliminate zero-length operands from Concat
    def has_zero_size(operand: ir.Value | None) -> bool:
        if operand is None:
            return False  # Invalid model
        if (shape := operand.shape) is None:
            return False
        try:
            # We have already checked that axis is an int value (!= None)
            dim_size = shape[axis]  # type: ignore[index]
        except IndexError:
            return False
        return dim_size == 0  # return False if symbolic or None or non-zero int value

    new_inputs = [x for x in inputs if not has_zero_size(x)]
    if len(new_inputs) != len(inputs):
        if new_inputs:
            # Remove zero-length operands from Concat
            logger.debug(
                "Concat: removing zero-length operand(s) %s => %s", inputs, new_inputs
            )
            return op.Concat(*new_inputs, axis=axis)
        elif inputs:
            # All operands are zero-length. Concat is a no-op, but we need to use one of the
            # inputs to get the other dimensions correct:
            logger.debug("Concat: removing all zero-length operands %s", inputs)
            return op.Identity(inputs[0])
        else:
            # No inputs: invalid model.
            return None

    # Track value of tensors that carry a shape value:

    # Check axis attribute is 0

    if axis != 0:
        return None
    shapes = [state.get_shape_value(input) for input in inputs]
    if any(shape is None for shape in shapes):
        return None
    concatenated = ir.Shape(dim for shape in shapes for dim in shape.dims)  # type: ignore[union-attr]
    output = node.outputs[0]
    if output is None:
        return None
    state.set_sym_value(output, concatenated)
    return None


@register("Dropout", version=(12, None))
def dropout(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace a Dropout by Identity when applicable."""

    def optimized_dropout():
        input = node.inputs[0]
        output = op.Identity(input)
        if len(node.outputs) == 1:
            return output
        else:
            true_tensor = ir.tensor([True])
            input_shape = op.Shape(input)
            mask = op.ConstantOfShape(input_shape, value=true_tensor)
            return output, mask

    inputs = node.inputs
    if (len(inputs) <= 2) or inputs[2] is None:
        # No training_mode specified:
        return optimized_dropout()
    if _get_bool_value(inputs[2]) is False:
        # training_mode is False: dropout is not applied.
        return optimized_dropout()
    ratio = _get_numpy_value(inputs[1])
    if ratio is None:
        return None
    if ratio.size != 1:  # Only scalar dropout ratio is supported.
        return None
    if ratio.item() == 0:
        # dropout ratio is 0: dropout is not applied.
        return optimized_dropout()
    return None


@register("Expand")
def expand(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Replace an Expand node by Identity when applicable."""
    if len(node.inputs) != 2:
        return None
    if (input := node.inputs[0]) is None:
        return None
    if (input_shape := input.shape) is None:
        # Input shape is not known.
        return None
    if (expanded_shape := _get_numpy_value(node.inputs[1])) is None:
        # Target shape is not known.
        expanded_sym_shape = state.get_shape_value(node.inputs[1])
        if expanded_sym_shape is None or not _same_shape(input_shape, expanded_sym_shape):
            return None
        return op.Identity(input)
    if expanded_shape.ndim != 1:
        # Target shape must be a 1D tensor. Erroneous model.
        return None
    if input_shape.dims == tuple(expanded_shape.tolist()):
        return op.Identity(input)
    return None


@register("ConcatFromSequence")
def concat_from_sequence(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input = node.inputs[0]
    inputs = state.get_sym_value(input)
    if inputs is None or any(x is None for x in inputs):
        return None
    new_axis = _get_int_attribute(node, "new_axis", 0)
    axis = _get_int_attribute(node, "axis", None)
    if axis is None:
        return None
    if input is not None and isinstance(inputs, list):
        if new_axis == 0:
            logger.debug("ConcatFromSequence => Concat: %s", [x.name for x in inputs])
            return op.Concat(*inputs, axis=axis)
        if new_axis == 1:
            # Unsqueeze the inputs with concat axis if new_axis is 1
            axis_value = op.Constant(value_int=axis)
            unsqueezed_inputs = []
            for node_input in inputs:
                unsqueezed_input = op.Unsqueeze(
                    node_input, axis_value, _outputs=[f"{node_input.name}_unsqueeze"]
                )
                unsqueezed_inputs.append(unsqueezed_input)
            # Send unsqueezed outputs to Concat
            logger.debug(
                "ConcatFromSequence => Concat %s", [x.name for x in unsqueezed_inputs]
            )
            return op.Concat(*unsqueezed_inputs, axis=axis)
    return None


@register("SplitToSequence")
def split_to_sequence(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    """Rewriting pattern.
    From

        splits = onnx::SplitToSequence(input, split, axis=axis)

    to

        split_0, split_1, ..., split_n = onnx::Split(input, split, axis=axis)
        splits = onnx::SequenceConstruct(split_0, split_1, ..., split_n)

    or

        split_0, split_1, ..., split_n = onnx::Split(input, axis=axis, num_outputs=n+1)
        splits = onnx::SequenceConstruct(split_0, split_1, ..., split_n)

    where number of output tensors in `splits` is statically known.
    onnx::SequenceConstruct will be further optimized away if possible, by its own designated evaluator.
    This allows downstream `SequenceAt` users to be replaced by `split_x` accordingly.
    """
    input = node.inputs[0]
    if len(node.inputs) == 1:
        # split is not provided
        return None
    split = node.inputs[1]
    output = node.outputs[0]

    if input is None or split is None or output is None:
        return None

    axis = _get_int_attribute(node, "axis", 0)
    if axis is None:
        return None
    shape = input.shape
    if shape is None:
        return None
    rank = len(shape)
    if axis < 0:
        axis = axis + rank
    if axis < 0 or axis >= rank:
        return None

    # NOTE: Split needs to either be a scalar or a 1-D tensor. We need to
    # calculate the number of outputs for Split.
    # If split is a scalar, we split into chunks of size 'split' if possible.
    #   * the split dimension size and split_value has to be known.
    # If split is a 1-D tensor, we split into 'size(split)' chunks
    #   * Get the size from split_value if it's numpy array.
    #   * Get the size from symbolic shape if split_value is not available.
    split_value = _get_numpy_value(split)
    split_shape = (
        split.shape.numpy() if split.shape is not None and split.shape.is_static() else None
    )

    # No information about split value or shape.
    if split_value is None and split_shape is None:
        return None

    if isinstance(split_shape, tuple) and len(split_shape) == 1:
        # If split_shape is known, we can use it to determine the number of outputs.
        split_dimension_size = split_shape[0]
        assert isinstance(split_dimension_size, int)
        num_outputs = split_dimension_size
        split_outputs = [f"{output.name}_split_{i}" for i in range(num_outputs)]
        split_values = op.Split(input, split, axis=axis, _outputs=split_outputs)
    elif split_value.ndim == 1:
        # split into 'size(split)' chunks
        num_outputs = split_value.size
        split_outputs = [f"{output.name}_split_{i}" for i in range(num_outputs)]
        split_values = op.Split(input, split, axis=axis, _outputs=split_outputs)
    elif split_value.ndim == 0:
        # split into chunks all of size 'split' if possible.
        split_dimension_size = shape[axis]
        if not isinstance(split_dimension_size, int):
            return None
        num_outputs = math.ceil(split_dimension_size / split_value.item())
        split_outputs = [f"{output.name}_split_{i}" for i in range(num_outputs)]
        split_values = op.Split(
            input, axis=axis, num_outputs=num_outputs, _outputs=split_outputs
        )
    else:
        return None

    # If Split returns a single value, we need to wrap it into a list.
    if isinstance(split_values, ir.Value):
        split_values = [split_values]

    keepdims = _get_int_attribute(node, "keepdims", 1)
    if keepdims is None:
        return None
    if keepdims == 0:
        # squeeze the split dimension if keepdims is 0
        axis_val = op.Constant(value_ints=[axis], _outputs=[f"{output.name}_axis"])
        squeezed_values = []
        for i in range(num_outputs):
            squeezed = op.Squeeze(
                split_values[i], axis_val, _outputs=[f"{split_outputs[i]}_squeeze"]
            )
            squeezed_values.append(squeezed)
        split_values = squeezed_values

    logger.debug("SplitToSequence => Split + SequenceConstruct")

    if isinstance(split_values, ir.Value):
        split_values = [split_values]
    return op.SequenceConstruct(*split_values)


@register("SequenceAt")
def sequence_at(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
    input = node.inputs[0]
    position = node.inputs[1]
    output = node.outputs[0]
    if input is not None and position is not None:
        input_vals = state.get_sym_value(input)
        position_val = _get_numpy_value(position)
        if isinstance(input_vals, list) and position_val is not None:
            if position_val.size != 1:
                return None
            position_val = position_val.item()
            try:
                result = input_vals[position_val]  # type: ignore[index]
            except IndexError:
                return None
            state.set_sym_value(output, result)
            logger.debug("SequenceAt %s => %s", input.name, result.name)
            return op.Identity(result)
    return None


def _merge_shapes(
    preferred_shape: ir.Shape | None, other_shape: ir.Shape | None
) -> ir.Shape | None:
    """Merge two shapes, preferring dimensions from preferred_shapes."""

    def merge_dims(dim1, dim2):
        if dim1 == dim2:
            return dim1
        if not isinstance(dim1, ir.SymbolicDim):
            return dim1  # Prefer int value over symbolic dim
        if not isinstance(dim2, ir.SymbolicDim):
            return dim2
        if dim1.value is None:
            return dim2
        return dim1

    if preferred_shape is None:
        return other_shape
    if other_shape is None:
        return preferred_shape
    if len(preferred_shape) != len(other_shape):
        raise ValueError(
            f"Shapes must have the same rank, got preferred_shape={preferred_shape}, other_shape={other_shape}"
        )
    return ir.Shape(
        [merge_dims(dim1, dim2) for dim1, dim2 in zip(preferred_shape, other_shape)]
    )


def _record_contributing_values(original_node: ir.Node, replacement: Replacement) -> None:
    """Record the set of original input values that contributed to the constant-folded outputs."""
    folded_from: set[str] = set()
    for input in original_node.inputs:
        if input is None:
            continue
        folded_from.update(input.meta.get(FOLDED_FROM_KEY, set()))
        assert input.name is not None
        folded_from.add(input.name)

    for new_output in replacement.new_outputs:
        if new_output is None:
            continue
        new_output.meta[FOLDED_FROM_KEY] = folded_from
        # Store the string representation of the set to metadata_props to persist it across serialization
        new_output.metadata_props[FOLDED_FROM_KEY] = repr(sorted(folded_from))


class FoldConstantsPass(ir.passes.InPlacePass):
    """A pass that folds constant expressions in the model.
    Attributes:
        shape_inference: Whether to perform shape inference.
        input_size_limit: Maximum size of input tensors to fold.
        output_size_limit: Maximum size of output tensors to fold.
        should_fold: An optional function that takes a node and returns True if
            the node should be considered for folding.
        The function should return True/False value to indicate if this particular
            node should be folded, or None to use the default folding rules.
    """

    def __init__(
        self,
        *,
        shape_inference: bool,
        input_size_limit: int,
        output_size_limit: int,
        should_fold: Callable[[ir.Node], bool | None] = lambda node: None,
    ) -> None:
        self.shape_inference = shape_inference
        self.input_size_limit = input_size_limit
        self.output_size_limit = output_size_limit
        self.should_fold = should_fold

        self._opset_imports: dict[str, int] = {}
        self._counts: dict[str, int] = {}
        self._sizes: dict[str, int] = {}
        self._modified: bool = False
        self._state = OptimizerState()
        self._reset()

    def _reset(self) -> None:
        """Reset internal states for a new run."""
        self._counts = {}
        self._sizes = {}
        self._modified = False
        self._state = OptimizerState()

    def _do_inference(self, node: ir.Node) -> None:
        output_types = {}

        # TODO: handle optional inputs
        def get_constant_value(x: ir.Value) -> onnx.TensorProto | None:
            value = _get_numpy_value(x, size_limit=20)
            if value is not None:
                assert x.const_value is not None
                return ir.serde.serialize_tensor(x.const_value)
            return None

        def get_type(value: ir.Value) -> onnx.TypeProto | None:
            if value.type is not None:
                type_proto = ir.serde.serialize_type(value.type)
                if value.shape is not None:
                    ir.serde.serialize_shape_into(type_proto, value.shape)
                return type_proto
            return None

        input_types = {x.name: get_type(x) for x in node.inputs if x is not None}
        input_data = {x.name: get_constant_value(x) for x in node.inputs if x is not None}
        input_data = {k: v for k, v in input_data.items() if v is not None}
        if any(t is None for t in input_types.values()):
            logger.debug(
                "Skipping shape inference for node %r due to missing input type.",
                node.name,
            )
        else:
            # TODO: pass in constant values, ir_version
            try:
                schema = onnx.defs.get_schema(
                    node.op_type, self._opset_imports[node.domain], node.domain
                )
                output_types = onnx.shape_inference.infer_node_outputs(
                    schema,
                    ir.serde.serialize_node(node),
                    input_types,  # type: ignore[arg-type]
                    input_data,  # type: ignore[arg-type]
                )
                for output in node.outputs:
                    if output.name in output_types:
                        inferred_type = output_types[output.name]
                        # TODO: merge types, check for conflicts
                        inferred_shape = ir.serde.deserialize_type_proto_for_shape(
                            inferred_type
                        )
                        # NOTE: forward shape inference
                        output.shape = _merge_shapes(output.shape, inferred_shape)
                        output.type = ir.serde.deserialize_type_proto_for_type(inferred_type)
            except Exception as e:
                logger.debug(
                    "Skipping shape inference for node %r due to exception: %s",
                    node,
                    e,
                )

    def _prepare_folded_tensor(
        self, node: ir.Node, output_name: str, output_array: np.ndarray | Any
    ) -> ir.Tensor | None:
        """
        Shared helper for constant/init creation:
        - Validates the folded Python value is a numpy ndarray.
        - Wraps it in an ir.Tensor and names it.
        - Applies output_size_limit logic with input-usage compensation.
        Returns the ir.Tensor or None if it should be skipped.
        """
        if not isinstance(output_array, np.ndarray):
            logger.info(
                "Skip storing constant folded value %s due to unsupported type %s.",
                output_name,
                type(output_array),
            )
            return None

        tensor = ir.tensor(output_array)
        tensor.name = output_name

        # Size gating (shared logic)
        if output_array.size > self.output_size_limit:
            removed_input_size = 0
            for input_val in node.inputs:
                if (input_val is not None) and (len(input_val.uses()) == 1):
                    input_array = _get_numpy_value(input_val)
                    if input_array is not None:
                        removed_input_size += input_array.size
            increased_size = output_array.size - removed_input_size
            if increased_size > 0:
                logger.info(
                    "Skip storing constant folded array %s due to large size %s.",
                    output_name,
                    output_array.size,
                )
                return None

        return tensor

    def new_constant(self, node: ir.Node, array: np.ndarray | Any) -> ir.Node | None:
        """Create a new Constant node with the given array as its value."""
        original_value = node.outputs[0]

        tensor = self._prepare_folded_tensor(node, original_value.name, array)
        if tensor is None:
            return None

        logger.debug(
            "New constant for value %s dtype: %s shape: %s",
            original_value.name,
            array.dtype,
            array.shape,
        )

        node = ir.Node("", "Constant", inputs=[], attributes=(ir.AttrTensor("value", tensor),))
        return node

    def new_initializer(self, node: ir.Node, array: np.ndarray | Any) -> ir.Value | None:
        """Create a new initializer value with the given array as its value."""
        original_value = node.outputs[0]

        tensor = self._prepare_folded_tensor(node, original_value.name, array)
        if tensor is None:
            return None

        initializer = ir.Value(
            name=original_value.name,
            type=ir.TensorType(ir.DataType(tensor.dtype)),
            shape=tensor.shape,  # type: ignore[arg-type]
            const_value=tensor,
        )

        logger.debug(
            "New Initializer for value %s dtype: %s shape: %s",
            original_value.name,
            array.dtype,
            array.shape,
        )

        return initializer

    def process_node(self, node: ir.Node, is_function: bool) -> Replacement | None:
        """Process a node and return a Replacement if the node can be replaced."""
        for i, value in enumerate(node.inputs):
            sym_value = self._state.get_sym_value(value)
            if isinstance(sym_value, ir.Value):
                logger.debug(
                    "Node [%s]: Replacing input %s with %s",
                    node.name,
                    value.name,  # type: ignore[union-attr]
                    sym_value.name,
                )
                node.replace_input_with(i, sym_value)
                self._modified = True
                # TODO(rama): consider merging type/other info from both values

        # Propagate const_value, and manually find out shape and type
        # to avoid potentially expensive shape inference on large tensors.
        if _is_onnx_op(node, "Constant"):
            _process_constant_node(node)
        # Do incremental shape inference
        elif self.shape_inference and not _is_control_flow_op(node):
            self._do_inference(node)

        if node.domain not in self._opset_imports:
            logger.debug(
                "Skipping constant folding for node %r due to missing opset import for domain %r.",
                node.name,
                node.domain,
            )
            return None

        version = self._opset_imports[node.domain]
        op_optimizers = registry.lookup_evaluators(node.domain, node.op_type, version)
        for optimizer in op_optimizers:
            assert optimizer
            context = RewriterContext()
            try:
                output = optimizer(node, context, self._state)
            except Exception as e:
                raise RuntimeError(
                    f"Error during constant folding for node {node.name!r} ({node.domain}::{node.op_type})"
                ) from e
            if output is not None:
                if isinstance(output, Replacement):
                    return output
                if isinstance(output, ir.Value):
                    output = [output]
                return Replacement(output, context.nodes)

        if _is_onnx_op(node, "Constant"):
            logger.debug("Skipping constant folding for Constant node %r", node.name)
            return None

        if _is_control_flow_op(node):
            logger.info(
                "Skipping constant folding for control flow op %r (%s::%s) because it is not supported yet",
                node.name,
                node.domain,
                node.op_type,
            )
            return None

        if _is_non_deterministic_op(node):
            logger.info(
                "Skipping constant folding for non-deterministic op %r (%s::%s)",
                node.name,
                node.domain,
                node.op_type,
            )
            return None

        if any(x.is_graph_input() for x in node.inputs if x is not None):
            logger.info(
                "Skipping constant folding for node %r because it is graph input to preserve graph signature",
                node.name,
            )
            return None

        # Ensure all node inputs are constants or initializers
        if any(x.const_value is None for x in node.inputs if x is not None):
            return None

        should_fold = self.should_fold(node)

        if should_fold is False:
            logger.info(
                "Skipping constant folding for node %r because should_fold returned False",
                node.name,
            )
            return None

        elif should_fold is None:
            # Use default rules to decide whether to fold the node:
            # - Nodes in the DEFAULT_CONSTANT_FOLD_BLACKLIST list are not folded
            # - If the any tensor input size exceeds the input_size_limit, skip folding the node
            for op_type in DEFAULT_CONSTANT_FOLD_BLACKLIST:
                if _is_onnx_op(node, op_type):
                    logger.info(
                        "Skipping constant folding for node %r because %s is preserved by default",
                        node.name,
                        op_type,
                    )
                    return None

            input_tensors = [x.const_value if x is not None else None for x in node.inputs]
            large_inputs = [
                tensor is not None and tensor.size > self.input_size_limit
                for tensor in input_tensors
            ]
            if any(large_inputs):
                # Decide whether to fold large constants
                assert len(node.inputs) == len(large_inputs)
                if (node.domain, node.op_type) in _DEFAULT_ALWAYS_FOLD_OPS and all(
                    len(input.consumers()) == 1 or (not is_large)
                    for input, is_large in zip(node.inputs, large_inputs)
                    if input is not None
                ):
                    # If the op is in _DEFAULT_ALWAYS_FOLD_OPS and all large inputs are used only by this node,
                    # we can still fold it even if the input size exceeds the limit
                    pass
                else:
                    # Skip folding large tensors
                    if logger.isEnabledFor(logging.INFO):
                        input_sizes = [
                            tensor.size for tensor in input_tensors if tensor is not None
                        ]
                        logger.info(
                            "Skipping constant folding for node %r due to large input sizes: %s",
                            node,
                            input_sizes,
                        )
                    return None
        else:
            logger.info(
                "Constant folding node %r because should_fold returned True",
                node.name,
            )

        input_values = [_get_numpy_value(x) for x in node.inputs]

        def convert(av):
            if av.type == ir.AttributeType.TENSOR:
                return ir.serde.serialize_tensor(av.value)
            return av.value

        # TODO(justinchuby): We should find a way to avoid serializing tensors every time we want to evaluate a node
        attr_values = {name: convert(attr) for name, attr in node.attributes.items()}
        outputs = _reference_evaluator.evaluate(
            node.domain, node.op_type, version, *input_values, **attr_values
        )

        if outputs is None:
            return None
        if len(node.outputs) == 1 and not isinstance(outputs, (tuple, list)):
            # We don't support initializers in functions, so we need to create Constant nodes
            if is_function:
                replacement = self.new_constant(node, outputs)
                if replacement is None:
                    return None
                return Replacement(replacement.outputs, [replacement])
            new_initializer_value = self.new_initializer(node, outputs)
            if new_initializer_value is None:
                return None
            # Add the new initializer to the graph
            assert node.graph is not None
            node.graph.register_initializer(new_initializer_value)
            return Replacement([new_initializer_value], [])
        else:
            logger.warning(
                "Skipping constant folding for op %s with multiple outputs.", node.op_type
            )
        return None

    def replace_node(
        self, node: ir.Node, replacement: Replacement, root: ir.Graph | ir.Function
    ) -> None:
        logger.debug("Replacing node: %s::%s %s", node.domain, node.op_type, node.name)

        # Record the names of the values that has contributed to the replacement
        _record_contributing_values(node, replacement)

        # Obtain the list of non-None inputs to the node before it is cleared by
        # replace_nodes_and_values to check for unused initializers later.
        node_inputs = [v for v in node.inputs if v is not None]

        ir.convenience.replace_nodes_and_values(
            root, node, [node], replacement.new_nodes, node.outputs, replacement.new_outputs
        )

        if isinstance(root, ir.Graph):
            # The old node should now be detached from the graph
            assert node.graph is None
            _clear_unused_initializers(node_inputs)

        self._modified = True

        # TODO: what about new opset_imports?
        # TODO: track statistics about replaced nodes and sizes of new constants

    def visit_attribute(self, attr: ir.Attr) -> None:
        if attr.is_ref():
            return
        if attr.type == ir.AttributeType.GRAPH:
            self.visit_graph(attr.as_graph())
        elif attr.type == ir.AttributeType.GRAPHS:
            for graph in attr.as_graphs():
                self.visit_graph(graph)

    def visit_node(self, node: ir.Node, root: ir.Graph | ir.Function) -> None:
        is_function = isinstance(root, ir.Function)
        replacement = self.process_node(node, is_function=is_function)
        if replacement is None:
            # No change. Process attributes.
            for attr in node.attributes.values():
                self.visit_attribute(attr)
            return
        else:
            self.replace_node(node, replacement, root)

    def visit_graph(self, graph: ir.Graph) -> None:
        for node in graph:
            self.visit_node(node, graph)

        # Replace outputs if output nodes can be folded. This are typically outputs from
        # Identity nodes
        for i, output in enumerate(graph.outputs):
            if output is None:
                continue
            sym_value = self._state.get_sym_value(output)
            if not isinstance(sym_value, ir.Value):
                # An output must be a Value
                continue
            if not _sym_value_can_replace_graph_output(graph, sym_value, output):
                continue
            # Rename sym_value to match the output name
            sym_value.name = output.name
            graph.outputs[i] = sym_value
            self._modified = True

    def visit_function(self, function: ir.Function) -> None:
        for node in function:
            self.visit_node(node, function)

    def call(self, model: ir.Model) -> FoldConstantsResult:
        self._reset()
        self._opset_imports = model.opset_imports
        self.visit_graph(model.graph)
        for function in model.functions.values():
            # TODO(rama): Should we specialize functions?
            self.visit_function(function)
        return FoldConstantsResult(model, self._modified, self._state.symbolic_value_map)


def _sym_value_can_replace_graph_output(
    graph: ir.Graph, sym_value: ir.Value, output: ir.Value
) -> bool:
    if (producer := sym_value.producer()) is None:
        # If the sym_value has no producer, it is some graph's input
        # ONNX does not allow a graph input to be a graph output
        return False
    if producer.graph is not graph:
        # The sym_value must be produced by a node in the graph to be an output of this graph
        return False
    if sym_value.is_graph_output():
        # If the sym_value is already an output of a graph, we cannot rename it
        # to this output name. Otherwise the graph output represented by sym_value
        # will lose its name.
        return False
    return True


def _clear_unused_initializers(values: Sequence[ir.Value]) -> None:
    # Detach all inputs to the node, then check for unused initializers
    for value in values:
        if value is None or not value.is_initializer():
            continue

        if (not value.uses()) and (not value.is_graph_output()):
            assert value.is_initializer()
            assert value.graph is not None
            assert value.name is not None
            value.graph.initializers.pop(value.name)


@dataclasses.dataclass
class FoldConstantsResult(ir.passes.PassResult):
    symbolic_value_map: dict[ir.Value, SymbolicValue]

    # Add conversion to bool for backward compatibility. The previously returned value
    # for the fold_constants method was a boolean indicating whether the model was modified.
    def __bool__(self) -> bool:
        return self.modified


def fold_constants(
    model: ir.Model,
    *,
    onnx_shape_inference: bool = False,
    input_size_limit: int = DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
    output_size_limit: int = DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
    should_fold: Callable[[ir.Node], bool | None] = lambda node: None,
) -> FoldConstantsResult:
    """
    Applies constant folding optimization to the model.

    Args:
        model: The ONNX model to optimize.
        onnx_shape_inference: Whether to enable ONNX shape inference during
            constant folding. Defaults to False.
        input_size_limit: The maximum size of input tensors
            that can be considered for constant folding. Defaults to
            `DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT`.
        output_size_limit: The maximum size of output tensors
            that can be stored after constant folding. Defaults to
            `DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT`.
        should_fold: An optional function that takes a node and returns True if
            the node should be considered for folding, False if it should not be folded,
            or None to use the default rules. Defaults to a function that always returns None.

    Returns:
        An instance of `FoldConstantsResult`.
    """
    folder_pass = FoldConstantsPass(
        shape_inference=onnx_shape_inference,
        input_size_limit=input_size_limit,
        output_size_limit=output_size_limit,
        should_fold=should_fold,
    )
    return folder_pass(model)