tract_core/ops/array/
scatter_nd.rs

1use crate::internal::*;
2use ndarray::*;
3
4#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Default)]
5pub enum ScatterReduction {
6    #[default]
7    None,
8    Add,
9    Mul,
10    Min,
11    Max,
12}
13
14impl ScatterReduction {
15    pub fn as_str(&self) -> &'static str {
16        match self {
17            ScatterReduction::None => "none",
18            ScatterReduction::Add => "add",
19            ScatterReduction::Mul => "mul",
20            ScatterReduction::Min => "min",
21            ScatterReduction::Max => "max",
22        }
23    }
24
25    pub fn parse(s: &str) -> TractResult<Self> {
26        Ok(match s {
27            "none" => ScatterReduction::None,
28            "add" => ScatterReduction::Add,
29            "mul" => ScatterReduction::Mul,
30            "min" => ScatterReduction::Min,
31            "max" => ScatterReduction::Max,
32            s => bail!("Unknown scatter reduction: {s}"),
33        })
34    }
35}
36
37#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
38pub struct ScatterNd {
39    pub reduction: ScatterReduction,
40}
41
42impl Op for ScatterNd {
43    fn name(&self) -> StaticName {
44        "ScatterNd".into()
45    }
46
47    op_as_typed_op!();
48}
49
50impl ScatterNd {
51    unsafe fn eval_t<T: Datum>(
52        data: &mut Tensor,
53        indices: &ArrayViewD<i64>,
54        updates: &TValue,
55    ) -> TractResult<()> {
56        let mut data = unsafe { data.to_array_view_mut_unchecked::<T>() };
57        let updates_plain = updates.try_as_plain()?;
58        let updates_view = unsafe { updates_plain.to_array_view_unchecked::<T>() };
59        for coords in tract_ndarray::indices(&indices.shape()[..indices.ndim() - 1]) {
60            let mut indices_into_data = indices.view();
61            let mut updates = updates_view.view();
62            for x in coords.slice() {
63                indices_into_data.index_axis_inplace(Axis(0), *x);
64                updates.index_axis_inplace(Axis(0), *x);
65            }
66            let mut data = data.view_mut();
67            for x in indices_into_data {
68                data.index_axis_inplace(Axis(0), *x as usize);
69            }
70            data.assign(&updates)
71        }
72        Ok(())
73    }
74
75    unsafe fn eval_t_reduce<T: Datum + PartialOrd + std::ops::AddAssign + std::ops::MulAssign>(
76        data: &mut Tensor,
77        indices: &ArrayViewD<i64>,
78        updates: &TValue,
79        reduction: ScatterReduction,
80    ) -> TractResult<()> {
81        let mut data = unsafe { data.to_array_view_mut_unchecked::<T>() };
82        let updates_plain = updates.try_as_plain()?;
83        let updates_view = unsafe { updates_plain.to_array_view_unchecked::<T>() };
84        for coords in tract_ndarray::indices(&indices.shape()[..indices.ndim() - 1]) {
85            let mut indices_into_data = indices.view();
86            let mut updates = updates_view.view();
87            for x in coords.slice() {
88                indices_into_data.index_axis_inplace(Axis(0), *x);
89                updates.index_axis_inplace(Axis(0), *x);
90            }
91            let mut data = data.view_mut();
92            for x in indices_into_data {
93                data.index_axis_inplace(Axis(0), *x as usize);
94            }
95            Zip::from(&mut data).and(&updates).for_each(|d, u| match reduction {
96                ScatterReduction::Add => *d += u.clone(),
97                ScatterReduction::Mul => *d *= u.clone(),
98                ScatterReduction::Min => {
99                    if u < d {
100                        *d = u.clone()
101                    }
102                }
103                ScatterReduction::Max => {
104                    if u > d {
105                        *d = u.clone()
106                    }
107                }
108                ScatterReduction::None => unreachable!(),
109            });
110        }
111        Ok(())
112    }
113}
114
115/// Locates the single axis along which `tuples` is the row-major enumeration of a
116/// contiguous block of `data_shape`, every other axis fully covered.
117///
118/// `tuples` is the flattened constant index tensor, `data_shape.len()` coordinates
119/// per tuple. Returns `(axis, start, len)` of the block, or `None` when the tuples
120/// are anything else: the comparison is exact and elementwise, so a match means the
121/// scatter writes exactly `data[.., start..start + len, ..]` once, in order.
122fn scattered_block(tuples: &[i64], data_shape: &[usize]) -> Option<(usize, usize, usize)> {
123    let rank = data_shape.len();
124    let count = tuples.len() / rank;
125    if count == 0 {
126        return None;
127    }
128    for axis in 0..rank {
129        let others: usize =
130            data_shape.iter().enumerate().filter(|(ax, _)| *ax != axis).map(|(_, d)| *d).product();
131        if others == 0 || !count.is_multiple_of(others) {
132            continue;
133        }
134        let len = count / others;
135        let Ok(start) = usize::try_from(tuples[axis]) else { continue };
136        if start + len > data_shape[axis] {
137            continue;
138        }
139        let mut block: TVec<usize> = data_shape.into();
140        block[axis] = len;
141        let canonical = tuples.chunks(rank).enumerate().all(|(pos, tuple)| {
142            let mut rest = pos;
143            (0..rank).rev().all(|ax| {
144                let coord = rest % block[ax];
145                rest /= block[ax];
146                tuple[ax] == (coord + if ax == axis { start } else { 0 }) as i64
147            })
148        });
149        if canonical {
150            return Some((axis, start, len));
151        }
152    }
153    None
154}
155
156impl TypedOp for ScatterNd {
157    as_op!();
158
159    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
160        Ok(tvec!(inputs[0].datum_type.fact(inputs[0].shape.to_tvec())))
161    }
162
163    /// Rewrites a constant-index block assignment into `Concat(axis, [Slice(data,
164    /// 0..start), updates, Slice(data, end..dim)])`, dropping empty slices.
165    ///
166    /// Fires only when the reduction is `None`, the indices are constant with a last
167    /// dimension equal to the data rank, the tuples are exactly the row-major
168    /// enumeration of a contiguous block along one axis with every other axis fully
169    /// covered, and `updates` matches that block in shape and datum type. Symbolic
170    /// shapes are declined: full coverage of an axis cannot be established against
171    /// constant indices. `TypedConcat::declutter` and `optim::slice` clean up behind
172    /// it.
173    fn declutter(
174        &self,
175        model: &TypedModel,
176        node: &TypedNode,
177    ) -> TractResult<Option<TypedModelPatch>> {
178        if self.reduction != ScatterReduction::None {
179            return Ok(None);
180        }
181        let (data, indices, updates) = args_3!(model.node_input_facts(node.id)?);
182        rule_if_some!(konst = &indices.konst);
183        rule_if_some!(data_shape = data.shape.as_concrete());
184        rule_if_some!(updates_shape = updates.shape.as_concrete());
185        if !data.is_plain()
186            || !updates.is_plain()
187            || data.datum_type != updates.datum_type
188            || konst.rank() < 2
189            || !konst.is_plain()
190            || *konst.shape().last().unwrap() != data_shape.len()
191        {
192            return Ok(None);
193        }
194        let tuples = konst.cast_to::<i64>()?;
195        let tuples = tuples.try_as_plain()?.as_slice::<i64>()?;
196        rule_if_some!((axis, start, len) = scattered_block(tuples, data_shape));
197        let mut block: TVec<usize> = data_shape.into();
198        block[axis] = len;
199        if updates_shape != &block[..] {
200            return Ok(None);
201        }
202
203        let mut patch = TypedModelPatch::new("ScatterNd as Slice/Concat");
204        let data_tap = patch.tap_model(model, node.inputs[0])?;
205        let mut parts = tvec!();
206        if start > 0 {
207            parts.push(
208                patch.wire_node(
209                    format!("{}.head", node.name),
210                    crate::ops::array::Slice::new(axis, 0, start),
211                    &[data_tap],
212                )?[0],
213            );
214        }
215        parts.push(patch.tap_model(model, node.inputs[2])?);
216        if start + len < data_shape[axis] {
217            parts.push(
218                patch.wire_node(
219                    format!("{}.tail", node.name),
220                    crate::ops::array::Slice::new(axis, start + len, data_shape[axis]),
221                    &[data_tap],
222                )?[0],
223            );
224        }
225        let wire = if parts.len() == 1 {
226            parts[0]
227        } else {
228            patch.wire_node(&node.name, crate::ops::array::TypedConcat::new(axis), &parts)?[0]
229        };
230        patch.shunt_outside(model, node.id.into(), wire)?;
231        Ok(Some(patch))
232    }
233}
234
235impl EvalOp for ScatterNd {
236    fn is_stateless(&self) -> bool {
237        true
238    }
239
240    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
241        let (data, indices, updates) = args_3!(inputs);
242        let indices = indices.cast_to::<i64>()?;
243        let indices = indices.to_plain_array_view::<i64>()?;
244        if data.datum_type() != updates.datum_type() {
245            bail!(
246                "Data and update must be of the same type, got {:?} and {:?}",
247                data.datum_type(),
248                updates.datum_type()
249            );
250        }
251        let mut data = data.into_tensor();
252        unsafe {
253            match self.reduction {
254                ScatterReduction::None => dispatch_datum_by_size!(
255                    Self::eval_t(data.datum_type())(&mut data, &indices, &updates)
256                )?,
257                reduction => dispatch_numbers!(Self::eval_t_reduce(data.datum_type())(
258                    &mut data, &indices, &updates, reduction
259                ))?,
260            }
261        }
262        Ok(tvec!(data.into_tvalue()))
263    }
264}
tract_core/ops/array/scatter_nd.rs

tract_core/ops/array/
scatter_nd.rs
