Crate onnx_runtime_shape_inference 

onnx-runtime-shape-inference

Symbolic shape inference over the onnx_runtime_ir::Graph IR.

This crate is the general, extensible successor to the bounded shape- inference stopgaps elsewhere in the runtime (the loader’s const-fold-lite pass and the session’s just-in-time data-dependent resolution). Its design mirrors the reference implementation justinchuby/onnx-shape-inference:

1. Extensible per-op registry keyed by (domain, op_type, opset) with range-based version matching (InferenceRegistry). Unregistered ops leave their outputs unresolved rather than failing.
2. Symbolic dimension arithmetic (DimExpr) — a small canonical integer polynomial that captures the affine/product forms the op set produces (d0*d1, d0+k, d0/k, reshape -1 cancellation), lowered back to IR Dims only when writing results.
3. Shape-DATA propagation (ShapeData) — tracks the known element values of the small integer tensors in Shape → Slice → Concat → Gather → Unsqueeze → Reshape chains, so computed shapes resolve without executing the graph. This is what lets transformer graphs infer statically.
4. Merge policies (MergePolicy) — Strict (concrete disagreements are errors) and Permissive (prefer the more specific dim and keep going; the robust default).

Usage

use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
let registry = InferenceRegistry::default_registry();
let opsets = graph.opset_imports.clone();
let report = registry
    .infer_graph(graph, &opsets, MergePolicy::Permissive)
    .expect("inference");
assert!(report.fully_resolved());

Single-node inference (for testing or custom passes) is available via InferenceRegistry::infer_node.

Design invariants

• Model-agnostic. Rules dispatch purely on (domain, op_type, opset) and tensor metadata — never on model names or op counts.
• The IR contract is not modified. Derived dimensions live in this crate’s DimExpr and are lowered to a fresh symbol when they cannot be expressed as an IR Dim.
• Permissive by default, never panics on unknown input. Errors are reserved for genuine contract violations (see ShapeInferError).

Control flow and the container-type limitation

Control-flow ops that carry subgraph bodies are inferred by propagating shapes through the body: If reconciles its two branch outputs, while Loop and Scan seed the body’s formal inputs from the node’s operands, infer the body, then map the body outputs back (stacking a trip-count / scan axis where the op requires one).

The Sequence family (SequenceEmpty/Construct/Insert/Erase/At/ Length/ConcatFromSequence/SplitToSequence), Optional (Optional/OptionalHasElement/OptionalGetElement), and Map ops need a container element type that a plain tensor TypeInfo cannot express. ValueType adds that additively: it wraps (never replaces) TypeInfo, so a value with no recorded ValueType is a plain tensor and the tensor-only path is byte-identical. The full Sequence family is registered, container types thread through the control-flow bodies (If/Loop/Scan/ SequenceMap) and across subgraph scope capture. The Optional and Map op handlers remain a smaller staged follow-up (the ValueType::Optional/Map representation already exists). See issues #355 and #449.

Re-exports

pub use context::InferenceContext;

pub use context::MergePolicy;

pub use context::NodeIo;

pub use context::SymbolInterner;

pub use context::TensorType;

pub use context::TypeInfo;

pub use context::TypedShape;

pub use context::ValueType;

pub use context::merge_shapes;

pub use dim_expr::DimExpr;

pub use shape_data::MAX_SHAPE_DATA_ELEMS;

pub use shape_data::ShapeData;

Modules

context

The inference context handed to each op rule, plus the supporting type model: TypeInfo, TypedShape, MergePolicy, and the SymbolInterner that lowers derived DimExprs back to IR Dims.

dim_expr

Symbolic dimension arithmetic: DimExpr.

shape_data

Shape-DATA propagation: ShapeData.

Structs

InferenceRegistry

A registry mapping (domain, op_type, opset) to an InferenceFn.

InferenceReport

A summary of what whole-graph inference resolved.

Enums

ShapeInferError

An error produced while inferring shapes.

Type Aliases

InferenceFn

An operator inference rule: reads inputs from the InferenceContext and sets its outputs’ types (and, where applicable, shape-data).