onnx_runtime_shape_inference/
lib.rs

1//! # `onnx-runtime-shape-inference`
2//!
3//! Symbolic shape inference over the [`onnx_runtime_ir::Graph`] IR.
4//!
5//! This crate is the general, extensible successor to the bounded shape-
6//! inference stopgaps elsewhere in the runtime (the loader's `const-fold-lite`
7//! pass and the session's just-in-time data-dependent resolution). Its design
8//! mirrors the reference implementation
9//! [`justinchuby/onnx-shape-inference`](https://github.com/justinchuby/onnx-shape-inference):
10//!
11//! 1. **Extensible per-op registry** keyed by `(domain, op_type, opset)` with
12//!    range-based version matching ([`InferenceRegistry`]). Unregistered ops
13//!    leave their outputs unresolved rather than failing.
14//! 2. **Symbolic dimension arithmetic** ([`DimExpr`]) — a small canonical
15//!    integer polynomial that captures the affine/product forms the op set
16//!    produces (`d0*d1`, `d0+k`, `d0/k`, reshape `-1` cancellation), lowered
17//!    back to IR [`Dim`](onnx_runtime_ir::Dim)s only when writing results.
18//! 3. **Shape-DATA propagation** ([`ShapeData`]) — tracks the known element
19//!    values of the small integer tensors in `Shape → Slice → Concat → Gather
20//!    → Unsqueeze → Reshape` chains, so computed shapes resolve without
21//!    executing the graph. This is what lets transformer graphs infer
22//!    statically.
23//! 4. **Merge policies** ([`MergePolicy`]) — [`Strict`](MergePolicy::Strict)
24//!    (concrete disagreements are errors) and
25//!    [`Permissive`](MergePolicy::Permissive) (prefer the more specific dim and
26//!    keep going; the robust default).
27//!
28//! ## Usage
29//!
30//! ```no_run
31//! use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
32//! # fn demo(graph: &mut onnx_runtime_ir::Graph) {
33//! let registry = InferenceRegistry::default_registry();
34//! let opsets = graph.opset_imports.clone();
35//! let report = registry
36//!     .infer_graph(graph, &opsets, MergePolicy::Permissive)
37//!     .expect("inference");
38//! assert!(report.fully_resolved());
39//! # }
40//! ```
41//!
42//! Single-node inference (for testing or custom passes) is available via
43//! [`InferenceRegistry::infer_node`].
44//!
45//! ## Design invariants
46//!
47//! * **Model-agnostic.** Rules dispatch purely on `(domain, op_type, opset)` and
48//!   tensor metadata — never on model names or op counts.
49//! * **The IR contract is not modified.** Derived dimensions live in this
50//!   crate's [`DimExpr`] and are lowered to a fresh symbol when they cannot be
51//!   expressed as an IR [`Dim`](onnx_runtime_ir::Dim).
52//! * **Permissive by default, never panics on unknown input.** Errors are
53//!   reserved for genuine contract violations (see [`ShapeInferError`]).
54//!
55//! ## Control flow and the container-type limitation
56//!
57//! Control-flow ops that carry subgraph bodies are inferred by propagating
58//! shapes *through* the body: `If` reconciles its two branch
59//! outputs, while `Loop` and `Scan` seed the body's formal inputs from the
60//! node's operands, infer the body, then map the body outputs back (stacking a
61//! trip-count / scan axis where the op requires one).
62//!
63//! The **Sequence** family (`SequenceEmpty`/`Construct`/`Insert`/`Erase`/`At`/
64//! `Length`/`ConcatFromSequence`/`SplitToSequence`), **Optional**
65//! (`Optional`/`OptionalHasElement`/`OptionalGetElement`), and **Map** ops need
66//! a *container element type* that a plain tensor [`TypeInfo`] cannot express.
67//! [`ValueType`] adds that additively: it wraps (never replaces) [`TypeInfo`],
68//! so a value with no recorded `ValueType` is a plain tensor and the tensor-only
69//! path is byte-identical. The full **Sequence** family is registered, container
70//! types thread through the control-flow bodies (`If`/`Loop`/`Scan`/
71//! `SequenceMap`) and across subgraph scope capture. The **Optional** and **Map**
72//! op handlers remain a smaller staged follow-up (the `ValueType::Optional`/`Map`
73//! representation already exists). See issues #355 and #449.
74
75#![forbid(unsafe_code)]
76
77pub mod context;
78pub mod dim_expr;
79mod error;
80mod handlers;
81mod infer;
82mod registry;
83mod report;
84pub mod shape_data;
85
86pub use context::{
87    InferenceContext, MergePolicy, NodeIo, SymbolInterner, TensorType, TypeInfo, TypedShape,
88    ValueType, merge_shapes,
89};
90pub use dim_expr::DimExpr;
91pub use error::ShapeInferError;
92pub use registry::{InferenceFn, InferenceRegistry};
93pub use report::InferenceReport;
94pub use shape_data::{MAX_SHAPE_DATA_ELEMS, ShapeData};
onnx_runtime_shape_inference/lib.rs

onnx_runtime_shape_inference/
lib.rs
