tract_hir/ops/array/
gather_nd.rs

1use crate::internal::*;
2pub use tract_core::ops::array::GatherNd;
3
4impl InferenceRulesOp for GatherNd {
5    fn rules<'r, 'p: 'r, 's: 'r>(
6        &'s self,
7        s: &mut Solver<'r>,
8        inputs: &'p [TensorProxy],
9        outputs: &'p [TensorProxy],
10    ) -> InferenceResult {
11        check_input_arity(inputs, 2)?;
12        check_output_arity(outputs, 1)?;
13        s.equals(&outputs[0].datum_type, &inputs[0].datum_type)?;
14        s.given(&inputs[1].rank, move |s, indices_rank| {
15            let indices_rank = indices_rank as usize;
16            for i in 0..(indices_rank - 1) {
17                s.equals(&outputs[0].shape[i], &inputs[1].shape[i])?;
18            }
19            s.given_2(
20                &inputs[1].shape[indices_rank - 1],
21                &inputs[1].rank,
22                move |s, n, input_rank| {
23                    if let Ok(n) = n.to_i64() {
24                        for i in 0..(input_rank - n) as usize {
25                            s.equals(&outputs[0].shape[indices_rank - 1 + i], &inputs[1].shape[i])?;
26                        }
27                    }
28                    Ok(())
29                },
30            )
31        })
32    }
33
34    as_op!();
35    to_typed!();
36}
tract_hir/ops/array/gather_nd.rs

tract_hir/ops/array/
gather_nd.rs
