[WIP] Create symbolic type/shape inference logic

pyproject.toml

@@ -21,7 +21,7 @@ classifiers = [
   "Programming Language :: Python :: 3.13",
   "License :: OSI Approved :: Apache Software License",
 ]
-dependencies = ["numpy", "onnx>=1.16", "typing_extensions>=4.10", "ml_dtypes"]
+dependencies = ["numpy", "onnx>=1.16", "typing_extensions>=4.10", "ml_dtypes", "sympy"]
 
 [project.urls]
 Homepage = "https://onnx.ai/ir-py"

src/onnx_ir/_core.py

@@ -45,6 +45,8 @@
 
 import ml_dtypes
 import numpy as np
+import sympy
+import sympy.utilities.misc
 from typing_extensions import TypeIs
 
 import onnx_ir
@@ -1115,13 +1117,14 @@
     It is immutable and can be compared or hashed.
     """
 
-    __slots__ = ("_value",)
+    __slots__ = ("_expr", "_value")
 
-    def __init__(self, value: str | None) -> None:
+    def __init__(self, value: str | None, /, expr: sympy.Expr | None = None) -> None:
         """Initialize a symbolic dimension.
 
         Args:
             value: The value of the dimension. It should not be an int.
+            expr: An optional sympy expression representing the dimension.
 
         Raises:
             TypeError: If value is an int.
@@ -1132,6 +1135,7 @@
                 "If you are creating a Shape, use int directly instead of SymbolicDim."
             )
         self._value = value
+        self._expr: sympy.Expr | None = expr
 
     def __eq__(self, other: object) -> bool:
         """Check equality with another SymbolicDim or string/None."""
@@ -1148,11 +1152,24 @@
         """The value of the symbolic dimension (string or None)."""
         return self._value
 
+    @property
+    def expr(self) -> sympy.Expr | None:
+        """The sympy expression representing the symbolic dimension."""
+        return self._expr
+
     def __str__(self) -> str:
-        return f"{self._value}"
+        if self._value is not None:
+            return str(self._value)
+        if self._expr is not None:
+            return str(self._expr)
+        return "?"
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self._value})"
+        if self._expr is not None:
+            expr_text = f", expr={self._expr!r}"
+        else:
+            expr_text = ""
+        return f"{self.__class__.__name__}({self._value}{expr_text})"
 
 
 def _is_int_compatible(value: object) -> TypeIs[SupportsInt]:
@@ -1190,10 +1207,16 @@
         return SymbolicDim(dim)
     if _is_int_compatible(dim):
         return int(dim)
+    if isinstance(dim, sympy.Expr):
+        # If the dimension is a sympy expression, we create a SymbolicDim with it
+        expr = sympy.sympify(dim)
+        if expr.is_integer:
+            return sympy.utilities.misc.as_int(expr)
+        return SymbolicDim(str(expr), expr=sympy.sympify(expr))
     if isinstance(dim, SymbolicDim):
         return dim
     raise TypeError(
-        f"Expected int, str, None or SymbolicDim, but value {dim!r} has type '{type(dim)}'"
+        f"Expected int, str, sympy.Expr, None or SymbolicDim, but value {dim!r} has type '{type(dim)}'"
     )
 
 
@@ -1334,7 +1357,9 @@
     def __getitem__(self, index):
         return tuple(self._dims)[index]
 
-    def __setitem__(self, index: int, value: int | SymbolicDim | str | None) -> None:
+    def __setitem__(
+        self, index: int, value: int | SymbolicDim | str | sympy.Expr | None
+    ) -> None:
         """Set the dimension at the index.
 
         Args:

src/onnx_ir/_shape_type_inference/__init__.py

@@ -0,0 +1,15 @@
+"""Symbolic shape and type inference for ONNX IR."""
+
+__all__ = [
+    "SymbolicInferenceEngine",
+    "InferenceError",
+    "NodeInferrer",
+    "InferenceResult",
+]
+
+
+from onnx_ir._shape_type_inference._common import InferenceResult, NodeInferrer
+from onnx_ir._shape_type_inference._engine import (
+    InferenceError,
+    SymbolicInferenceEngine,
+)

src/onnx_ir/_shape_type_inference/_common.py

@@ -0,0 +1,159 @@
+"""Symbolic shape inference for ONNX IR."""
+
+from __future__ import annotations
+
+import abc
+import dataclasses
+import functools
+from collections.abc import Collection, Sequence
+from typing import Any, Callable
+
+import sympy
+
+import onnx_ir as ir
+
+
+MAX_SUPPORTED_OPSET = 23
+
+
+def get_expr(shape: ir.Shape, index: int) -> sympy.Expr:
+    """Get the expression or value at a specific index in the shape.
+
+    Args:
+        shape: The shape to get the expression from.
+        index: The index of the dimension to get.
+
+    Returns:
+        The expression or value at the specified index.
+    """
+    dim = shape[index]
+    if isinstance(dim, ir.SymbolicDim):
+        if dim.expr is not None:
+            return dim.expr
+        if dim.value is None:
+            return sympy.Symbol("__unknown__")
+        return sympy.Symbol(dim.value)
+    return sympy.Integer(dim)
+
+
+@dataclasses.dataclass
+class InferenceResult:
+    values: Sequence[ir.Value] | None = None
+    failure: str | None = None
+
+
+class NodeInferrer(abc.ABC):
+    """Base class for node inferrers.
+
+    This class provides a common interface for all node inferrers.
+    """
+
+    def __init__(self, op_type: str, opsets: Collection[int], domain: str = "") -> None:
+        """Initialize the node inferrer.
+
+        Args:
+            op_type: The type of the operation.
+            opsets: A collection of ONNX opset versions supported by this inferrer.
+            domain: The domain of the operation, default is an empty string.
+        """
+        self.op_type = op_type
+        self.opsets = opsets
+        self.domain = domain
+
+    def __repr__(self) -> str:
+        """Return a string representation of the node inferrer."""
+        return f"{self.__class__.__name__}(op_type={self.op_type}, opsets={self.opsets}, domain={self.domain})"
+
+    @abc.abstractmethod
+    def infer(self, node: ir.Node) -> InferenceResult:
+        """Infer the shape for the node.
+
+        Args:
+            node: The ONNX node to infer the type and shape for.
+
+        Returns:
+            A sequence of ONNX values containing the inferred shapes.
+        """
+        raise NotImplementedError
+
+
+def requires_non_none_inputs(
+    count: int, /
+) -> Callable[
+    [Callable[[Any, ir.Node], InferenceResult]], Callable[[Any, ir.Node], InferenceResult]
+]:
+    """Ensure that the node has a specific number of non-None inputs.
+
+    Args:
+        count: The exact number of non-None inputs required for the node.
+
+    Returns:
+        A decorator that checks the number of inputs and their non-None status.
+    """
+
+    def decorator(
+        func: Callable[[Any, ir.Node], InferenceResult],
+    ) -> Callable[[Any, ir.Node], InferenceResult]:
+        @functools.wraps(func)
+        def wrapper(self, node: ir.Node) -> InferenceResult:
+            if len(node.inputs) != count:
+                return InferenceResult(
+                    failure=f"[{node.op_type} must have {count} inputs, got {len(node.inputs)}."
+                )
+            for i, inp in enumerate(node.inputs):
+                if inp is None:
+                    return InferenceResult(failure=f"{node.op_type} input {i} cannot be None.")
+            return func(self, node)
+
+        return wrapper
+
+    return decorator
+
+
+def requires_outputs(
+    count: int, /
+) -> Callable[
+    [Callable[[Any, ir.Node], InferenceResult]], Callable[[Any, ir.Node], InferenceResult]
+]:
+    """Ensure that the node has a specific number of outputs.
+
+    Args:
+        count: The exact number of outputs required for the node.
+
+    Returns:
+        A decorator that checks the number of outputs.
+    """
+
+    def decorator(
+        func: Callable[[Any, ir.Node], InferenceResult],
+    ) -> Callable[[Any, ir.Node], InferenceResult]:
+        @functools.wraps(func)
+        def wrapper(self, node: ir.Node) -> InferenceResult:
+            if len(node.outputs) != count:
+                return InferenceResult(
+                    failure=f"[{node.op_type} must have {count} outputs, got {len(node.outputs)}."
+                )
+            return func(self, node)
+
+        return wrapper
+
+    return decorator
+
+
+def inclusive_range(start_or_end: int = 0, end: int | None = None) -> range:
+    """Create an inclusive range from start to end with a given step.
+
+    Args:
+        start_or_end: The starting value of the range.
+        end: The ending value of the range (inclusive).
+
+    Returns:
+        A range object that includes both start and end.
+    """
+    if end is None:
+        end = start_or_end
+        start = 0
+    else:
+        start = start_or_end
+
+    return range(start, end + 1)