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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

//! Value-selection rules: `ArgMax`, `ArgMin`, `TopK`, and `NonZero`.

use onnx_runtime_ir::{Attribute, DataType};

use crate::context::InferenceContext;
use crate::dim_expr::DimExpr;
use crate::error::ShapeInferError;
use crate::handlers::norm_axis;
use crate::registry::InferenceRegistry;

/// `ArgMax`/`ArgMin`: replace or remove the selected axis and return int64.
fn arg_reduce(ctx: &mut InferenceContext) -> Result<(), ShapeInferError> {
    let Some(input) = ctx.input_shape(0).map(<[DimExpr]>::to_vec) else {
        return Ok(());
    };
    if input.is_empty() {
        return Ok(());
    }
    let axis = ctx
        .node
        .attr("axis")
        .and_then(Attribute::as_int)
        .unwrap_or(0);
    let axis = norm_axis(axis, input.len());
    let keepdims = ctx
        .node
        .attr("keepdims")
        .and_then(Attribute::as_int)
        .unwrap_or(1)
        != 0;
    let mut output = input;
    if keepdims {
        output[axis] = DimExpr::constant(1);
    } else {
        output.remove(axis);
    }
    ctx.set_output(0, DataType::Int64, output);
    Ok(())
}

/// Read a statically-known scalar integer operand.
fn scalar_int(ctx: &InferenceContext, input: usize) -> Option<i64> {
    let data = ctx.input_shape_data(input)?;
    if data.elems.len() != 1 {
        return None;
    }
    data.elems.first()?.as_const()
}

/// `TopK`: values retain the data dtype and indices are int64, with the selected
/// axis replaced by `K`.
fn top_k_v1(ctx: &mut InferenceContext) -> Result<(), ShapeInferError> {
    let k = ctx.node.attr("k").and_then(Attribute::as_int);
    top_k(ctx, k)
}

fn top_k_v10(ctx: &mut InferenceContext) -> Result<(), ShapeInferError> {
    top_k(ctx, scalar_int(ctx, 1))
}

fn top_k(ctx: &mut InferenceContext, k: Option<i64>) -> Result<(), ShapeInferError> {
    let Some(input) = ctx.input_shape(0).map(<[DimExpr]>::to_vec) else {
        return Ok(());
    };
    let Some(dtype) = ctx.input_dtype(0) else {
        return Ok(());
    };
    if input.is_empty() {
        return Ok(());
    }
    let axis = ctx
        .node
        .attr("axis")
        .and_then(Attribute::as_int)
        .unwrap_or(-1);
    let axis = norm_axis(axis, input.len());
    let mut output = input;
    output[axis] = k
        .filter(|&k| k >= 0)
        .map(DimExpr::constant)
        .unwrap_or_else(|| ctx.fresh_dim());
    ctx.set_output(0, dtype, output.clone());
    ctx.set_output(1, DataType::Int64, output);
    Ok(())
}

/// `NonZero`: its second dimension is data-dependent, while the first is the
/// statically-known rank of the input.
fn non_zero(ctx: &mut InferenceContext) -> Result<(), ShapeInferError> {
    let Some(rank) = ctx.input_rank(0) else {
        return Ok(());
    };
    let nnz = ctx.fresh_dim();
    ctx.set_output(
        0,
        DataType::Int64,
        vec![DimExpr::constant(rank as i64), nnz],
    );
    Ok(())
}

/// Register selection-family rules.
pub fn register(reg: &mut InferenceRegistry) {
    for op in ["ArgMax", "ArgMin"] {
        reg.register("", op, 1, arg_reduce);
        reg.register("", op, 11, arg_reduce);
        reg.register("", op, 12, arg_reduce);
        reg.register("", op, 13, arg_reduce);
    }
    reg.register("", "TopK", 1, top_k_v1);
    reg.register("", "TopK", 10, top_k_v10);
    reg.register("", "TopK", 11, top_k_v10);
    reg.register("", "NonZero", 9, non_zero);
    reg.register("", "NonZero", 13, non_zero);
}