tract_core/ops/array/
gather_nd.rs

1use crate::internal::*;
2use tract_ndarray::prelude::*;
3
4#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
5pub struct GatherNd {
6    pub batch_dims: usize,
7}
8
9impl GatherNd {
10    fn compute_shape<D: DimLike>(
11        &self,
12        data_shape: &[D],
13        indices_shape: &[D],
14    ) -> TractResult<TVec<D>> {
15        let mut shape: TVec<D> = indices_shape.into();
16        let n = shape.pop().unwrap().to_usize()?;
17        shape.extend(data_shape[n + self.batch_dims..].iter().cloned());
18        Ok(shape)
19    }
20
21    unsafe fn eval_t<T: Datum>(
22        &self,
23        output: &mut Tensor,
24        data: &Tensor,
25        indices: &ArrayViewD<i32>,
26    ) -> TractResult<()> {
27        let batch_dims = self.batch_dims;
28        assert_eq!(output.shape()[..batch_dims], data.shape()[..batch_dims]);
29        assert_eq!(output.shape()[..batch_dims], indices.shape()[..batch_dims]);
30        let batch_size = data.shape().iter().take(batch_dims).product();
31        let n = indices.shape()[indices.ndim() - 1];
32
33        let remaining = indices.shape().iter().skip(batch_dims).rev().skip(1).product();
34        let indices_shape_op = tvec!(batch_size, remaining, n);
35        let reshaped_indices: ArrayViewD<i32> =
36            indices.view().into_shape_with_order(&*indices_shape_op).unwrap();
37
38        let mut data_shape_op: TVec<usize> =
39            data.shape().iter().skip(batch_dims).copied().collect();
40        data_shape_op.insert(0, batch_size);
41        let data_plain = data.try_as_plain()?;
42        let reshaped_data = unsafe {
43            data_plain
44                .to_array_view_unchecked::<T>()
45                .into_shape_with_order(&*data_shape_op)
46                .unwrap()
47        };
48
49        let mut output_shape_op: TVec<usize> =
50            data.shape().iter().skip(n + batch_dims).copied().collect();
51        output_shape_op.insert(0, batch_size * remaining);
52        let mut output_plain = output.try_as_plain_mut()?;
53        let mut output = unsafe {
54            output_plain
55                .to_array_view_mut_unchecked::<T>()
56                .into_shape_with_order(&*output_shape_op)
57                .unwrap()
58        };
59
60        for b in 0..batch_size {
61            let mut i = reshaped_data.view();
62            i.index_axis_inplace(Axis(0), b);
63            let mut coords = reshaped_indices.view();
64            coords.index_axis_inplace(Axis(0), b);
65
66            for ix in 0..remaining {
67                let mut coords = coords.view();
68                coords.index_axis_inplace(Axis(0), ix);
69
70                let mut i = i.view();
71                for x in coords {
72                    i.index_axis_inplace(Axis(0), *x as usize);
73                }
74
75                let mut o = output.view_mut();
76                o.index_axis_inplace(Axis(0), b * remaining + ix);
77                o.assign(&i);
78            }
79        }
80        Ok(())
81    }
82}
83
84impl Op for GatherNd {
85    fn name(&self) -> StaticName {
86        "GatherNd".into()
87    }
88
89    op_as_typed_op!();
90}
91
92impl EvalOp for GatherNd {
93    op_out_of_plan!();
94
95    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
96        let (data, indices) = args_2!(inputs);
97        let shape = self.compute_shape(data.shape(), indices.shape())?;
98        let indices = indices.cast_to::<i32>()?;
99        let indices = indices.to_plain_array_view::<i32>()?;
100        unsafe {
101            let mut output = Tensor::uninitialized_dt(data.datum_type(), &shape)?;
102            dispatch_datum_by_size!(Self::eval_t(data.datum_type())(
103                self,
104                &mut output,
105                &data,
106                &indices
107            ))?;
108            Ok(tvec!(output.into_tvalue()))
109        }
110    }
111}
112
113impl TypedOp for GatherNd {
114    as_op!();
115
116    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
117        let shape = self.compute_shape(&inputs[0].shape.to_tvec(), &inputs[1].shape.to_tvec())?;
118        Ok(tvec!(inputs[0].datum_type.fact(&shape)))
119    }
120
121    fn declutter(
122        &self,
123        model: &TypedModel,
124        node: &TypedNode,
125    ) -> TractResult<Option<TypedModelPatch>> {
126        if let Some(indices) = &model.outlet_fact(node.inputs[1])?.konst
127            && indices.rank() == 2
128            && indices.shape()[0] == 1
129        {
130            let mut patch = TypedModelPatch::default();
131            let mut wire = patch.tap_model(model, node.inputs[0])?;
132            for (axis, &i) in
133                indices.cast_to::<i32>()?.try_as_plain()?.as_slice::<i32>()?.iter().enumerate()
134            {
135                wire = patch.wire_node(
136                    format!("{}-slice-axis-{}", node.name, axis),
137                    crate::ops::array::Slice::new(axis, i as usize, (i + 1) as usize),
138                    &[wire],
139                )?[0];
140            }
141            for i in (0..indices.shape()[1]).rev() {
142                wire = patch.wire_node(
143                    format!("{}-remove_axis_{}", node.name, i),
144                    crate::ops::change_axes::AxisOp::Rm(i),
145                    &[wire],
146                )?[0];
147            }
148            wire = patch.wire_node(
149                format!("{}-add_axis", node.name),
150                crate::ops::change_axes::AxisOp::Add(0),
151                &[wire],
152            )?[0];
153            patch.shunt_outside(model, node.id.into(), wire)?;
154            return Ok(Some(patch));
155        }
156        Ok(None)
157    }
158}
tract_core/ops/array/gather_nd.rs

tract_core/ops/array/
gather_nd.rs
