oxionnx 0.1.2

Pure Rust ONNX inference engine — zero C/C++ dependencies
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

//! Element-wise and variadic broadcast helpers.

use crate::graph::Node;
use crate::tensor::Tensor;
use std::collections::HashMap;

use crate::optimizer::shape_inference::get_input_shape;

/// Broadcast all variadic inputs together.
pub(super) fn infer_variadic_broadcast(
    node: &Node,
    known: &HashMap<String, Vec<usize>>,
) -> Option<Vec<Vec<usize>>> {
    if node.inputs.is_empty() {
        return None;
    }
    let mut result = get_input_shape(node, 0, known)?;
    for i in 1..node.inputs.len() {
        let shape = get_input_shape(node, i, known)?;
        result = Tensor::broadcast_shape(&result, &shape).ok()?;
    }
    Some(vec![result])
}