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#![forbid(unsafe_code)]
56
57pub mod context;
58pub mod dim_expr;
59mod error;
60mod handlers;
61mod infer;
62mod registry;
63mod report;
64pub mod shape_data;
65
66pub use context::{
67    InferenceContext, MergePolicy, NodeIo, SymbolInterner, TypeInfo, TypedShape, merge_shapes,
68};
69pub use dim_expr::DimExpr;
70pub use error::ShapeInferError;
71pub use registry::{InferenceFn, InferenceRegistry};
72pub use report::InferenceReport;
73pub use shape_data::{MAX_SHAPE_DATA_ELEMS, ShapeData};
onnx_runtime_shape_inference/lib.rs

onnx_runtime_shape_inference/
lib.rs
