Source code for tensorplay.export.dynamic_shapes

"""Dynamic dimension declarations and input-shape utilities."""

from __future__ import annotations

import dataclasses
import inspect
import re
import sys
from collections.abc import Callable, Mapping
from enum import Enum, auto
from collections.abc import Iterator
from typing import Any

__all__ = [
    "AdditionalInputs",
    "Constraint",
    "ConstraintsExceededError",
    "Dim",
    "DerivedDim",
    "ShapesCollection",
    "dims",
    "refine_dynamic_shapes_from_suggested_fixes",
]



[docs]
class ConstraintsExceededError(RuntimeError):
    """A runtime input violated the declared dynamic-shape contract.

    Raised by the assertions inserted into captured graphs and by export-time
    validation when an example input falls outside a declared range.  It is a
    ``RuntimeError`` so callers written against plain runtime failures keep
    working.
    """



class _DimHintType(Enum):
    AUTO = auto()
    STATIC = auto()
    DYNAMIC = auto()


@dataclasses.dataclass(frozen=True)
class _DimHint:
    type: _DimHintType
    min: int | None = None
    max: int | None = None
    _factory: bool = True

    @staticmethod
    def AUTO() -> "_DimHint":
        return _DimHint(_DimHintType.AUTO)

    @staticmethod
    def STATIC() -> "_DimHint":
        return _DimHint(_DimHintType.STATIC)

    @staticmethod
    def DYNAMIC() -> "_DimHint":
        return _DimHint(_DimHintType.DYNAMIC)

    def __call__(self, min: int | None = None, max: int | None = None) -> "_DimHint":
        if not self._factory:
            raise TypeError(f"{type(self).__name__!s} object is not callable")
        if min is not None and (type(min) is not int or min < 0):
            raise ValueError(f"min must be a non-negative integer, got {min!r}")
        if max is not None and (type(max) is not int or max < 0):
            raise ValueError(f"max must be a non-negative integer, got {max!r}")
        if min is not None and max is not None and min > max:
            raise ValueError(f"min must be no greater than max, got {min} and {max}")
        return _DimHint(self.type, min, max, False)

    def __repr__(self) -> str:
        values = [self.type.name]
        if self.min is not None:
            values.append(f"min={self.min}")
        if self.max is not None:
            values.append(f"max={self.max}")
        return f"DimHint({', '.join(values)})"



[docs]
class Dim:
    """A named symbolic dimension with an optional finite range."""

    AUTO = _DimHint.AUTO()
    STATIC = _DimHint.STATIC()
    DYNAMIC = _DimHint.DYNAMIC()

    def __init__(
        self,
        name: str,
        *,
        min: int | None = None,
        max: int | None = None,
    ) -> None:
        if not isinstance(name, str) or not name.isidentifier():
            raise ValueError(f"dimension name must be an identifier, got {name!r}")
        lower = 0 if min is None else min
        upper = max
        if type(lower) is not int or lower < 0:
            raise ValueError(f"min must be a non-negative integer, got {min!r}")
        if upper is not None and (type(upper) is not int or upper < lower):
            raise ValueError(f"max must be an integer no less than min, got {max!r}")
        self.__name__ = name
        self.min = lower
        self.max = upper

    @property
    def name(self) -> str:
        return self.__name__

    def _derive(self, scale: int, offset: int) -> "_DerivedDim":
        return _DerivedDim(_linear_name(self.__name__, scale, offset), self, scale, offset)

    def __add__(self, other: Any) -> "_DerivedDim":
        if type(other) is not int:
            raise NotImplementedError("dimension addition requires an integer")
        return self._derive(1, other)

    def __radd__(self, other: Any) -> "_DerivedDim":
        return self + other

    def __sub__(self, other: Any) -> "_DerivedDim":
        if type(other) is not int:
            raise NotImplementedError("dimension subtraction requires an integer")
        return self._derive(1, -other)

    def __rsub__(self, other: Any) -> "_DerivedDim":
        raise NotImplementedError("a dimension cannot be negated")

    def __mul__(self, other: Any) -> "_DerivedDim":
        if type(other) is not int or other <= 0:
            raise NotImplementedError("dimension multiplication requires a positive integer")
        return self._derive(other, 0)

    def __rmul__(self, other: Any) -> "_DerivedDim":
        return self * other

    def __repr__(self) -> str:
        bounds = []
        if self.min != 0:
            bounds.append(f"min={self.min}")
        if self.max is not None:
            bounds.append(f"max={self.max}")
        suffix = f", {', '.join(bounds)}" if bounds else ""
        return f"Dim({self.__name__!r}{suffix})"



def _linear_name(root: str, scale: int, offset: int) -> str:
    """Render ``scale * root + offset`` in canonical expression form."""

    if scale == 1 and offset == 0:
        return root
    base = root if scale == 1 else f"{scale}*{root}"
    if offset == 0:
        return base
    sign = "+" if offset > 0 else "-"
    return f"{base} {sign} {abs(offset)}"


class _StaticDim(Dim):
    """Dimension pinned to one concrete size by an integer specification."""

    def __init__(self, value: int) -> None:
        if type(value) is not int or isinstance(value, bool) or value < 0:
            raise ValueError(f"static dimension requires a non-negative int, got {value!r}")
        self.__name__ = str(value)
        self.value = value

    @property
    def name(self) -> str:
        return self.__name__

    @property
    def min(self) -> int:
        return self.value

    @property
    def max(self) -> int:
        return self.value

    def __repr__(self) -> str:
        return self.__name__


class _DerivedDim(Dim):
    """A positive linear expression ``scale * root + offset`` over one base dimension."""

    def __init__(self, name: str, root: Dim, scale: int, offset: int) -> None:
        if isinstance(root, _DerivedDim):
            scale, offset, root = scale * root.scale, scale * root.offset + offset, root.root
        if scale <= 0:
            raise NotImplementedError("derived dimensions require a positive scale")
        self.__name__ = name
        self.root = root
        self.scale = scale
        self.offset = offset

    @property
    def name(self) -> str:
        return self.__name__

    def _evaluate(self, value: int) -> int:
        return self.scale * value + self.offset

    @property
    def min(self) -> int:
        value = self._evaluate(self.root.min)
        if value < 0:
            raise ValueError(
                f"derived dimension {self.__name__!r} has a negative lower bound; "
                f"specify a larger min for the root {self.root.__name__!r}"
            )
        return value

    @property
    def max(self) -> int | None:
        if self.root.max is None:
            return None
        value = self._evaluate(self.root.max)
        if value > sys.maxsize - 1:
            raise ValueError(f"derived dimension {self.__name__!r} exceeds the integer range")
        return value

    def _derive(self, scale: int, offset: int) -> "_DerivedDim":
        return _DerivedDim(
            _linear_name(self.root.__name__, self.scale * scale, self._compose_offset(offset)),
            self.root,
            self.scale * scale,
            scale * self.offset + offset,
        )

    def _compose_offset(self, offset: int) -> int:
        return self.offset + offset

    def __repr__(self) -> str:
        return self.__name__


DerivedDim = _DerivedDim
"""Public alias for a linear expression ``scale * root + offset`` over one
base dimension."""



[docs]
def dims(*names: str, min: int | None = None, max: int | None = None) -> tuple[Dim, ...]:
    """Construct several named dimensions with shared bounds."""

    return tuple(Dim(name, min=min, max=max) for name in names)




[docs]
@dataclasses.dataclass(frozen=True)
class Constraint:
    """A range restriction attached to one input dimension.

    ``name`` ties the constraint to a named :class:`Dim` shared across
    inputs (equalities are implied).  ``root``/``scale``/``offset`` describe a
    derived dimension whose size equals ``scale * root_size + offset``.
    Constraints without a ``name`` come from dim hints or static entries.
    """

    source: Any
    dim: int
    name: str | None = None
    min: int | None = None
    max: int | None = None
    warn_only: bool = False
    root: str | None = None
    scale: int = 1
    offset: int = 0

    def _replace_range(self, *, min: int | None = None, max: int | None = None) -> "Constraint":
        lower = self.min if min is None else min
        upper = self.max if max is None else max
        if upper is not None and lower is not None and lower > upper:
            raise ValueError("constraint range is empty")
        return dataclasses.replace(self, min=lower, max=upper)

    def __ge__(self, value: int) -> "Constraint":
        return self._replace_range(min=value)

    def __gt__(self, value: int) -> "Constraint":
        return self._replace_range(min=value + 1)

    def __le__(self, value: int) -> "Constraint":
        return self._replace_range(max=value)

    def __lt__(self, value: int) -> "Constraint":
        return self._replace_range(max=value - 1)

    def __bool__(self) -> bool:
        raise TypeError("a dimension constraint cannot be used as a boolean")

    @property
    def serializable_spec(self) -> dict[str, Any]:
        spec = {"source": self.source, "dim": self.dim, "min": self.min, "max": self.max}
        if self.name is not None:
            spec["name"] = self.name
        if self.root is not None:
            spec["root"] = self.root
            spec["scale"] = self.scale
            spec["offset"] = self.offset
        return spec



class _IntWrapper:
    def __init__(self, value: int) -> None:
        self.value = value


def _is_shape_value(value: Any) -> bool:
    return hasattr(value, "shape") or isinstance(value, _IntWrapper)


def _map_tree(value: Any, fn: Callable[[Any], Any]) -> Any:
    if isinstance(value, dict):
        return {key: _map_tree(item, fn) for key, item in value.items()}
    if isinstance(value, list):
        return [_map_tree(item, fn) for item in value]
    if isinstance(value, tuple):
        return type(value)(_map_tree(item, fn) for item in value)
    return fn(value)


def _combine_args(model: Any, args: tuple[Any, ...], kwargs: Mapping[str, Any] | None) -> Any:
    callable_obj = getattr(model, "forward", model)
    signature = inspect.signature(callable_obj)
    bound = signature.bind_partial(*args, **dict(kwargs or {}))
    bound.apply_defaults()
    return dict(bound.arguments)



[docs]
class ShapesCollection:
    """Associate shape specifications with tensor objects by identity."""

    def __init__(self) -> None:
        self._shapes: dict[int, Any] = {}

    def __setitem__(self, value: Any, shape: Any) -> None:
        if not _is_shape_value(value):
            raise TypeError(f"cannot assign a shape to {type(value).__name__}")
        key = id(value)
        previous = self._shapes.get(key, inspect.Parameter.empty)
        if previous is not inspect.Parameter.empty and previous != shape:
            raise ValueError(f"shape already assigned as {previous!r}")
        self._shapes[key] = shape

    def __getitem__(self, value: Any) -> Any:
        return self._shapes.setdefault(id(value), {})

    def __len__(self) -> int:
        return len(self._shapes)

    def dynamic_shapes(
        self, model: Any, args: tuple[Any, ...], kwargs: Mapping[str, Any] | None = None
    ) -> Any:
        seen: set[int] = set()

        def find(value: Any) -> Any:
            key = id(value)
            if key in self._shapes:
                seen.add(key)
                return self._shapes[key]
            return None

        result = _map_tree(_combine_args(model, args, kwargs), find)
        missing = set(self._shapes) - seen
        if missing:
            raise ValueError("some assigned shape values were not found in the inputs")
        return result



def _shape_snapshot(value: Any) -> Any:
    if hasattr(value, "shape"):
        try:
            return tuple(int(item) for item in value.shape)
        except Exception:
            return tuple(value.shape)
    if isinstance(value, dict):
        return {key: _shape_snapshot(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_shape_snapshot(item) for item in value]
    if isinstance(value, tuple):
        return tuple(_shape_snapshot(item) for item in value)
    return value


def _mark_dynamism(value: Any, *others: Any) -> Any:
    if others and any(type(value) is not type(other) for other in others):
        raise ValueError("additional inputs have incompatible value types")
    if isinstance(value, int) and not isinstance(value, bool):
        return None if all(value == other for other in others) else Dim.DYNAMIC
    if any(value != other for other in others):
        raise ValueError("additional inputs have incompatible static values")
    return None



[docs]
class AdditionalInputs:
    """Infer dynamic shape markers from representative input sets."""

    def __init__(self) -> None:
        self._examples: list[tuple[tuple[Any, ...], dict[str, Any] | None]] = []

    def add(self, args: tuple[Any, ...], kwargs: Mapping[str, Any] | None = None) -> None:
        if type(args) is not tuple:
            raise TypeError("representative args must be a tuple")
        if kwargs is not None and type(kwargs) is not dict:
            raise TypeError("representative kwargs must be a dict or None")
        self._examples.append((args, None if kwargs is None else dict(kwargs)))

    def dynamic_shapes(
        self, model: Any, args: tuple[Any, ...], kwargs: Mapping[str, Any] | None = None
    ) -> Any:
        snapshots = [
            _shape_snapshot(_combine_args(model, current_args, current_kwargs))
            for current_args, current_kwargs in [(args, kwargs), *self._examples]
        ]
        if len(snapshots) == 1:
            return _map_tree(snapshots[0], lambda value: None)

        def merge(values: list[Any]) -> Any:
            first = values[0]
            if isinstance(first, dict):
                if not all(isinstance(item, dict) and item.keys() == first.keys() for item in values):
                    raise ValueError("additional inputs have incompatible mappings")
                return {key: merge([item[key] for item in values]) for key in first}
            if isinstance(first, (list, tuple)):
                if not all(type(item) is type(first) and len(item) == len(first) for item in values):
                    raise ValueError("additional inputs have incompatible sequences")
                result = [merge([item[index] for item in values]) for index in range(len(first))]
                return type(first)(result)
            return _mark_dynamism(first, *values[1:])

        return merge(snapshots)

    def verify(self, program: Any) -> None:
        module = program.module()
        for args, kwargs in self._examples:
            module(*args, **(kwargs or {}))



def _replace_dim(value: Any, replacements: Mapping[str, Any]) -> Any:
    if isinstance(value, Dim):
        return replacements.get(value.name, value)
    if isinstance(value, dict):
        return {key: _replace_dim(item, replacements) for key, item in value.items()}
    if isinstance(value, list):
        return [_replace_dim(item, replacements) for item in value]
    if isinstance(value, tuple):
        return tuple(_replace_dim(item, replacements) for item in value)
    return value


def _collect_named_dims(value: Any, found: dict[str, Dim] | None = None) -> dict[str, Dim]:
    found = {} if found is None else found
    if isinstance(value, _DerivedDim):
        found.setdefault(value.__name__, value)
        return _collect_named_dims(value.root, found)
    if isinstance(value, Dim):
        found.setdefault(value.__name__, value)
    elif isinstance(value, dict):
        for item in value.values():
            _collect_named_dims(item, found)
    elif isinstance(value, (list, tuple)):
        for item in value:
            _collect_named_dims(item, found)
    return found



[docs]
def refine_dynamic_shapes_from_suggested_fixes(message: str, dynamic_shapes: Any) -> Any:
    """Apply suggested fixes (range refinements, specializations, relations).

    Supported fix forms::

        name = Dim('name', min=..., max=...)  # refine a range
        name = 4                              # specialize to a constant
        dy = dx + 1                           # tie a dim to another with a relation
        dy = 2*dx                             # positive multiple of another dim

    ``dx`` must name a dimension already present in ``dynamic_shapes`` or be
    defined by an earlier fix line.
    """

    if not isinstance(message, str):
        raise TypeError("message must be a string")
    section = message.split("Suggested fixes:", 1)
    if len(section) != 2:
        raise ValueError("suggested fixes were not found")

    known_dims = _collect_named_dims(dynamic_shapes)
    fixes: list[tuple[str, str]] = []
    for line in section[1].splitlines():
        line = line.split("#", 1)[0].strip()
        match = re.match(r"([A-Za-z_]\w*)\s*=\s*(.+)", line)
        if match:
            fixes.append((match.group(1), match.group(2).strip()))
    if not fixes:
        raise ValueError("no supported shape fixes were found")

    resolved: dict[str, Any] = {}

    def resolve(expression: str) -> Any:
        if expression in resolved:
            return resolved[expression]
        dim_match = re.match(
            r"Dim\(['\"]([A-Za-z_]\w*)['\"](?:,\s*min=(\d+))?(?:,\s*max=(\d+))?\)",
            expression,
        )
        if dim_match:
            value = Dim(
                dim_match.group(1),
                min=int(dim_match.group(2)) if dim_match.group(2) else None,
                max=int(dim_match.group(3)) if dim_match.group(3) else None,
            )
            resolved[expression] = value
            return value
        if re.fullmatch(r"\d+", expression):
            value = int(expression)
            resolved[expression] = value
            return value
        linear = re.match(
            r"(?:(\d+)\s*\*\s*)?([A-Za-z_]\w*)(?:\s*([+-])\s*(\d+))?",
            expression,
        )
        if linear and linear.group(0).strip() == expression.strip():
            scale = int(linear.group(1)) if linear.group(1) else 1
            offset = int(linear.group(4) or 0)
            if linear.group(3) == "-":
                offset = -offset
            root_name = linear.group(2)
            if root_name in resolved and isinstance(resolved[root_name], Dim):
                root = resolved[root_name]
            elif root_name in known_dims:
                root = known_dims[root_name]
            else:
                raise ValueError(
                    f"fix references unknown dimension {root_name!r}; it must appear "
                    f"in dynamic_shapes or in an earlier suggested fix"
                )
            if scale == 1 and offset == 0:
                value = root
            elif offset:
                value = root * scale + offset
            else:
                value = root * scale
            resolved[expression] = value
            return value
        raise ValueError(f"unsupported suggested fix expression: {expression!r}")

    replacements: dict[str, Any] = {}
    for name, expression in fixes:
        value = resolve(expression)
        resolved[name] = value
        if value is not None:
            replacements[name] = value
    return _replace_dim(dynamic_shapes, replacements)



def _iter_pairs(values: Any, specs: Any, path: tuple[Any, ...] = ()) -> Iterator[tuple[tuple[Any, ...], Any, Any]]:
    """Walk an input tree and a specification tree in lockstep."""

    if isinstance(values, dict) and isinstance(specs, dict):
        for key in values:
            yield from _iter_pairs(values[key], specs.get(key), path + (key,))
    elif isinstance(values, (tuple, list)) and isinstance(specs, (tuple, list)):
        if len(specs) != len(values):
            raise ValueError(
                f"dynamic shape specification at inputs[{_render_path(path)}] has "
                f"{len(specs)} entries but the input has {len(values)}"
            )
        for index, (item, spec) in enumerate(zip(values, specs)):
            yield from _iter_pairs(item, spec, path + (index,))
    else:
        yield path, values, specs


def _render_path(path: tuple[Any, ...]) -> str:
    rendered = ""
    for item in path:
        rendered = f"{rendered}[{item!r}]" if isinstance(item, str) else f"{rendered}[{item}]"
    return rendered or "[]"


def _tensor_size(value: Any) -> tuple[int, ...] | None:
    shape = getattr(value, "shape", None)
    if shape is None or isinstance(value, _IntWrapper):
        return None
    try:
        return tuple(int(item) for item in shape)
    except TypeError:
        return None


def _suggest_fix(name: str, dim: Dim | None, size: int) -> str:
    if dim is None:
        return f"    {name} = {size}  # specialize to a constant"
    return f"    {name} = {dim!r}"


def _check_dynamic_shapes(
    combined_args: Any,
    dynamic_shapes: Any,
) -> None:
    """Validate a dynamic_shapes specification against the example inputs."""

    if not dynamic_shapes:
        return
    for path, value, spec in _iter_pairs(combined_args, dynamic_shapes):
        size = _tensor_size(value)
        if size is None:
            if spec is not None and not isinstance(spec, _DimHint):
                raise ValueError(
                    f"cannot associate shape {spec!r} at dynamic_shapes{_render_path(path)} "
                    f"with non-tensor input of type {type(value).__name__}"
                )
            continue
        if spec is None:
            continue
        entries = (
            list(spec.items()) if isinstance(spec, dict) else list(enumerate(spec))
        )
        for index, dim in entries:
            if dim is None:
                continue
            if not isinstance(index, int) or index < 0:
                raise TypeError(f"dimension index must be a non-negative int, got {index!r}")
            if index >= len(size):
                raise ValueError(
                    f"dimension index {index} at dynamic_shapes{_render_path(path)} is out "
                    f"of range for a {len(size)}-dimensional input"
                )
            if type(dim) is int and not isinstance(dim, bool):
                if size[index] != dim:
                    raise ValueError(
                        f"input size {size[index]} of dimension {index} at "
                        f"dynamic_shapes{_render_path(path)} does not match the expected "
                        f"static size {dim}"
                    )
                continue
            if isinstance(dim, _DimHint):
                if dim.min is not None and size[index] < dim.min:
                    raise ValueError(
                        f"input size {size[index]} of dimension {index} at "
                        f"dynamic_shapes{_render_path(path)} violates min={dim.min}"
                    )
                if dim.max is not None and size[index] > dim.max:
                    raise ValueError(
                        f"input size {size[index]} of dimension {index} at "
                        f"dynamic_shapes{_render_path(path)} violates max={dim.max}"
                    )
            elif isinstance(dim, _StaticDim):
                if size[index] != dim.value:
                    raise ValueError(
                        f"input size {size[index]} of dimension {index} at "
                        f"dynamic_shapes{_render_path(path)} does not match the expected "
                        f"static size {dim.value}"
                    )
            elif isinstance(dim, Dim):
                if size[index] < dim.min or (dim.max is not None and size[index] > dim.max):
                    lower = f"min={dim.min}" if dim.min else ""
                    upper = f"max={dim.max}" if dim.max is not None else ""
                    bound = ", ".join(item for item in (lower, upper) if item)
                    # widen to include the observed size: the repaired range
                    # covers the example input, keeping the dim dynamic
                    suggested = Dim(
                        dim.__name__,
                        min=min(dim.min, size[index]),
                        max=None if dim.max is None else max(dim.max, size[index]),
                    )
                    raise ValueError(
                        f"input size {size[index]} of dimension {index} at "
                        f"dynamic_shapes{_render_path(path)} violates {dim.__name__}"
                        f"({bound}). Suggested fixes:\n"
                        f"{_suggest_fix(dim.__name__, suggested, size[index])}"
                    )


def _process_dynamic_shapes(
    combined_args: Any,
    dynamic_shapes: Any,
) -> list[Constraint]:
    """Turn a validated dynamic_shapes specification into dimension constraints."""

    if not dynamic_shapes:
        return []
    _check_dynamic_shapes(combined_args, dynamic_shapes)
    constraints: list[Constraint] = []
    for path, value, spec in _iter_pairs(combined_args, dynamic_shapes):
        size = _tensor_size(value)
        if size is None or spec is None:
            continue
        entries = (
            list(spec.items()) if isinstance(spec, dict) else list(enumerate(spec))
        )
        for index, dim in entries:
            if dim is None:
                continue
            if isinstance(dim, _StaticDim):
                constraints.append(
                    Constraint(value, index, name=None, min=dim.value, max=dim.value)
                )
            elif isinstance(dim, _DerivedDim):
                constraints.append(
                    Constraint(
                        value,
                        index,
                        name=dim.__name__,
                        min=dim.min,
                        max=dim.max,
                        root=dim.root.__name__,
                        scale=dim.scale,
                        offset=dim.offset,
                    )
                )
            elif isinstance(dim, Dim):
                constraints.append(
                    Constraint(value, index, name=dim.__name__, min=dim.min, max=dim.max)
                )
            elif isinstance(dim, _DimHint):
                relaxed = dim.type is not _DimHintType.STATIC
                constraints.append(
                    Constraint(
                        value,
                        index,
                        name=None,
                        min=dim.min,
                        max=dim.max,
                        warn_only=relaxed,
                    )
                )
    return constraints


def _constraint_program(
    constraints: list[Constraint],
) -> tuple[list[Constraint], dict[str, dict[str, int | None]]]:
    """Group constraints into runtime assertions and per-name range bounds."""

    asserts: list[Constraint] = []
    ranges: dict[str, dict[str, int | None]] = {}
    anchors: dict[str, tuple[Any, int]] = {}
    for constraint in constraints:
        if constraint.root is not None:
            asserts.append(constraint)
            ranges.setdefault(
                constraint.name, {"min": constraint.min, "max": constraint.max}
            )
            continue
        if constraint.name is not None:
            if constraint.name in ranges:
                known = ranges[constraint.name]
                if known["min"] != constraint.min or known["max"] != constraint.max:
                    raise ValueError(
                        f"found conflicting definitions for symbolic dimension "
                        f"{constraint.name!r}: {known} and "
                        f"{{'min': {constraint.min}, 'max': {constraint.max}}}"
                    )
                if not constraint.warn_only:
                    asserts.append(constraint)  # equality with the anchor
            else:
                ranges[constraint.name] = {"min": constraint.min, "max": constraint.max}
                if not constraint.warn_only:
                    asserts.append(constraint)
        elif not constraint.warn_only:
            asserts.append(constraint)
    return asserts, ranges