Skip to content

shape_propagation

Classes

ShapeRuleResult

Base type for shape propagation rule outcomes.

NotApplicable dataclass 

NotApplicable()

Rule does not own this operator case and generic fallback may run.

Updated dataclass 

Updated(shapes: dict[str, tuple[int, ...]])

Rule produced concrete output shape updates.

NoUpdate dataclass 

NoUpdate(reason: str = '')

Rule owns this operator case and no shape update is required.

Blocked dataclass 

Blocked(reason: str)

Rule owns this operator case and cannot safely infer an update.

ShapeInferenceResult dataclass 

ShapeInferenceResult(changed_tensor_ids: set[str] = set(), confirmed_tensor_ids: set[str] = set(), unresolved_tensor_ids: set[str] = set(), diagnostics: list[str] = list())

Shape propagation outcome.

Attributes:

  • changed_tensor_ids (set[str]) –

    Tensor IDs whose shapes changed during propagation.

  • confirmed_tensor_ids (set[str]) –

    Tensor IDs whose shapes a rule computed and verified, whether or not the stored shape changed. A dynamic-signature tensor whose inferred shape already matched is resolved, not unresolved.

  • unresolved_tensor_ids (set[str]) –

    Dynamic tensor IDs that remain unresolved.

  • diagnostics (list[str]) –

    Human-readable diagnostics describing inference updates and unresolved tensors. Skip diagnostics (rule/fallback skipped) are only included when DEBUG logging is enabled to avoid unbounded growth on large graphs.

Functions

propagate_shapes 

propagate_shapes(model: AirModel, *, strict: bool = False) -> ShapeInferenceResult

Infer and propagate tensor shapes over the graph until fixed-point.

Parameters:

  • model

    (AirModel) –

    AIR model to update in place.

  • strict

    (bool, default: False ) –

    If True, raise when unresolved dynamic tensors remain.

Returns:

Raises:

  • ValueError –

    If strict is True and unresolved dynamic tensors remain.

  • TypeError –

    If an operator's options type does not match its shape rule's contract, or a rule returns an unknown result type.