rten 0.25.0

Machine learning runtime
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

//! Operator which computes a tensor shape by evaluating symbolic expressions.

use rten_base::num::AsUsize;
use rten_shape_inference::{SymExpr, SymbolMap};
use rten_tensor::Tensor;
use rten_tensor::prelude::*;

use crate::infer_shapes::InferShapes;
use crate::operator::{
    IntoOpResult, OpError, OpRunContext, Operator, OutputList, OutputType, OutputTypeList,
    OutputTypesContext,
};
use crate::value::{DataType, ValueType};

/// Specifies the source for a variable used in symbolic expressions.
#[derive(Clone, Debug, PartialEq)]
pub struct SymbolInfo {
    /// Name of the symbol in symbolic expressions.
    pub name: String,

    /// Input index.
    pub input: u32,

    /// Axis index.
    pub axis: u32,
}

/// Compute a tensor shape from a combination of static values and the dynamic
/// shapes of inputs.
#[derive(Debug)]
pub struct ComputeShape {
    /// Specifies how to map dimension sizes of inputs to symbols used by the
    /// `shape` field.
    pub symbols: Vec<SymbolInfo>,

    /// Specifies the symbolic expression to evaluate for each position in the
    /// output vector.
    pub shape: Vec<SymExpr>,
}

impl Operator for ComputeShape {
    fn name(&self) -> &str {
        "ComputeShape"
    }

    fn max_inputs(&self) -> Option<usize> {
        None
    }

    fn run(&self, ctx: &OpRunContext) -> Result<OutputList, OpError> {
        let inputs = ctx.inputs();

        let symbols = self
            .symbols
            .iter()
            .map(|sym| {
                let input = inputs.require(sym.input.as_usize())?;
                if input.ndim() <= sym.axis.as_usize() {
                    return Err(OpError::InvalidValue("Axis invalid for input shape"));
                }
                let size = input.size(sym.axis.as_usize()) as i32;
                Ok((sym.name.as_str(), size))
            })
            .collect::<Result<Vec<_>, _>>()?;
        let symbols = SymbolMap::new(&symbols);

        let output = self
            .shape
            .iter()
            .map(|expr| {
                expr.eval(&symbols)
                    .map_err(|_| OpError::InvalidValue("Failed to evaluate symbolic shape"))
            })
            .collect::<Result<Vec<i32>, _>>()?;

        Tensor::from(output).into_op_result()
    }

    fn output_types(&self, _ctx: &OutputTypesContext) -> Option<OutputTypeList> {
        Some([OutputType::Fixed(ValueType::Tensor(DataType::Int32))].into())
    }

    fn as_infer_shapes(&self) -> Option<&dyn InferShapes> {
        None
    }
}

#[cfg(test)]
mod tests {
    use rten_shape_inference::{SymExpr, Symbol};
    use rten_tensor::prelude::*;
    use rten_tensor::{NdTensor, NdTensorView};

    use super::{ComputeShape, SymbolInfo};
    use crate::operator::{OpError, OperatorExt};

    #[test]
    fn test_compute_shape() {
        // Valid input with a mix of static and dynamic output values.
        let input_a = NdTensor::<f32, _>::zeros([2, 4, 8]);
        let input_b = NdTensor::<f32, _>::zeros([24]);

        let op = ComputeShape {
            symbols: [
                SymbolInfo {
                    name: "x".to_string(),
                    input: 0,
                    axis: 1,
                },
                SymbolInfo {
                    name: "y".to_string(),
                    input: 1,
                    axis: 0,
                },
            ]
            .to_vec(),
            shape: [
                SymExpr::Value(3),
                SymExpr::Var(
                    Symbol {
                        name: "x".to_string(),
                        positive: true,
                        synthetic: false,
                    }
                    .into(),
                ),
                SymExpr::Value(5),
                SymExpr::Var(
                    Symbol {
                        name: "y".to_string(),
                        positive: true,
                        synthetic: false,
                    }
                    .into(),
                ),
            ]
            .into(),
        };
        let result: NdTensor<i32, 1> = op.run_simple((input_a.view(), input_b.view())).unwrap();

        assert_eq!(result, NdTensorView::from(&[3, 4, 5, 24]));

        // Dynamic input with invalid input index.
        let op = ComputeShape {
            symbols: [SymbolInfo {
                name: "x".to_string(),
                input: 1,
                axis: 0,
            }]
            .into(),
            shape: Vec::new(),
        };
        let result: Result<NdTensor<i32, 1>, _> = op.run_simple(input_a.view());
        assert_eq!(result.err().unwrap(), OpError::MissingInputs);

        // Dynamic input with invalid axis.
        let op = ComputeShape {
            symbols: [SymbolInfo {
                name: "x".to_string(),
                input: 0,
                axis: 3,
            }]
            .into(),
            shape: Vec::new(),
        };
        let result: Result<NdTensor<i32, 1>, _> = op.run_simple(input_a.view());
        assert_eq!(
            result.err().unwrap(),
            OpError::InvalidValue("Axis invalid for input shape")
        );
    }
}