onnx-runtime-shape-inference 0.1.0-dev.2

Symbolic shape inference for the ORT 2.0 runtime: an extensible, opset-aware per-op registry with symbolic dimension arithmetic and shape-data propagation over onnx-runtime-ir graphs
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

//! Built-in operator inference rules.
//!
//! Each rule is a small pure function over an [`InferenceContext`]. Rules are
//! grouped by family and wired into the [`InferenceRegistry`] by
//! [`register_all`]. The Phase-1 coverage targets the transformer + CNN op set
//! needed to fully infer BERT-class graphs without executing them; see the
//! crate docs for the full list.

use crate::registry::InferenceRegistry;

mod data_ops;
mod elementwise;
mod linalg;
mod movement;
mod norm;
mod pooling;
mod selection;
mod sequence;

/// Normalise an ONNX axis (which may be negative) into `0..rank`.
///
/// Negative axes count from the end. The result is clamped into `0..=rank-1`
/// for indexing an existing dimension (`rank` must be ≥ 1).
pub(crate) fn norm_axis(axis: i64, rank: usize) -> usize {
    let r = rank as i64;
    let a = if axis < 0 { axis + r } else { axis };
    a.clamp(0, r.saturating_sub(1)) as usize
}

/// Populate `registry` with every built-in rule.
pub fn register_all(registry: &mut InferenceRegistry) {
    elementwise::register(registry);
    linalg::register(registry);
    norm::register(registry);
    movement::register(registry);
    data_ops::register(registry);
    pooling::register(registry);
    selection::register(registry);
    sequence::register(registry);
}