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_ram()?;
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                if *x < 0 || *x as usize >= data.shape()[0] {
69                    anyhow::bail!(
70                        "ScatterND index {x} out of bounds for data shape {:?} (indices tensor {:?})",
71                        data.shape(),
72                        indices
73                    );
74                }
75                data.index_axis_inplace(Axis(0), *x as usize);
76            }
77            data.assign(&updates)
78        }
79        Ok(())
80    }
81
82    unsafe fn eval_t_reduce<T: Datum + PartialOrd + std::ops::AddAssign + std::ops::MulAssign>(
83        data: &mut Tensor,
84        indices: &ArrayViewD<i64>,
85        updates: &TValue,
86        reduction: ScatterReduction,
87    ) -> TractResult<()> {
88        let mut data = unsafe { data.to_array_view_mut_unchecked::<T>() };
89        let updates_plain = updates.try_as_plain_ram()?;
90        let updates_view = unsafe { updates_plain.to_array_view_unchecked::<T>() };
91        for coords in tract_ndarray::indices(&indices.shape()[..indices.ndim() - 1]) {
92            let mut indices_into_data = indices.view();
93            let mut updates = updates_view.view();
94            for x in coords.slice() {
95                indices_into_data.index_axis_inplace(Axis(0), *x);
96                updates.index_axis_inplace(Axis(0), *x);
97            }
98            let mut data = data.view_mut();
99            for x in indices_into_data {
100                data.index_axis_inplace(Axis(0), *x as usize);
101            }
102            Zip::from(&mut data).and(&updates).for_each(|d, u| match reduction {
103                ScatterReduction::Add => *d += u.clone(),
104                ScatterReduction::Mul => *d *= u.clone(),
105                ScatterReduction::Min => {
106                    if u < d {
107                        *d = u.clone()
108                    }
109                }
110                ScatterReduction::Max => {
111                    if u > d {
112                        *d = u.clone()
113                    }
114                }
115                ScatterReduction::None => unreachable!(),
116            });
117        }
118        Ok(())
119    }
120}
121
122/// Locates the single axis along which `tuples` is the row-major enumeration of a
123/// contiguous block of `data_shape`, every other axis fully covered.
124///
125/// `tuples` is the flattened constant index tensor, `data_shape.len()` coordinates
126/// per tuple. Returns `(axis, start, len)` of the block, or `None` when the tuples
127/// are anything else: the comparison is exact and elementwise, so a match means the
128/// scatter writes exactly `data[.., start..start + len, ..]` once, in order.
129fn scattered_block(tuples: &[i64], data_shape: &[usize]) -> Option<(usize, usize, usize)> {
130    let rank = data_shape.len();
131    let count = tuples.len() / rank;
132    if count == 0 {
133        return None;
134    }
135    for axis in 0..rank {
136        let others: usize =
137            data_shape.iter().enumerate().filter(|(ax, _)| *ax != axis).map(|(_, d)| *d).product();
138        if others == 0 || !count.is_multiple_of(others) {
139            continue;
140        }
141        let len = count / others;
142        let Ok(start) = usize::try_from(tuples[axis]) else { continue };
143        if start + len > data_shape[axis] {
144            continue;
145        }
146        let mut block: TVec<usize> = data_shape.into();
147        block[axis] = len;
148        let canonical = tuples.chunks(rank).enumerate().all(|(pos, tuple)| {
149            let mut rest = pos;
150            (0..rank).rev().all(|ax| {
151                let coord = rest % block[ax];
152                rest /= block[ax];
153                tuple[ax] == (coord + if ax == axis { start } else { 0 }) as i64
154            })
155        });
156        if canonical {
157            return Some((axis, start, len));
158        }
159    }
160    None
161}
162
163impl TypedOp for ScatterNd {
164    as_op!();
165
166    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
167        Ok(tvec!(inputs[0].datum_type.fact(inputs[0].shape.to_tvec())))
168    }
169
170    /// Rewrites a constant-index block assignment into `Concat(axis, [Slice(data,
171    /// 0..start), updates, Slice(data, end..dim)])`, dropping empty slices.
172    ///
173    /// Fires only when the reduction is `None`, the indices are constant with a last
174    /// dimension equal to the data rank, the tuples are exactly the row-major
175    /// enumeration of a contiguous block along one axis with every other axis fully
176    /// covered, and `updates` matches that block in shape and datum type. Symbolic
177    /// shapes are declined: full coverage of an axis cannot be established against
178    /// constant indices. `TypedConcat::declutter` and `optim::slice` clean up behind
179    /// it.
180    fn declutter(
181        &self,
182        model: &TypedModel,
183        node: &TypedNode,
184    ) -> TractResult<Option<TypedModelPatch>> {
185        rule_if!(self.reduction == ScatterReduction::None);
186        let (data, indices, updates) = args_3!(model.node_input_facts(node.id)?);
187        rule_if_some!(konst = &indices.konst);
188        rule_if_some!(data_shape = data.shape.as_concrete());
189        rule_if_some!(updates_shape = updates.shape.as_concrete());
190        rule_if!(data.is_plain() && updates.is_plain());
191        rule_if!(data.datum_type == updates.datum_type);
192        rule_if!(konst.rank() >= 2 && konst.is_plain_ram());
193        rule_if!(*konst.shape().last().unwrap() == data_shape.len());
194        let tuples = konst.cast_to::<i64>()?;
195        let tuples = tuples.try_as_plain_ram()?.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        rule_if!(updates_shape == &block[..]);
200
201        let mut patch = TypedModelPatch::new("ScatterNd as Slice/Concat");
202        let data_tap = patch.tap_model(model, node.inputs[0])?;
203        let mut parts = tvec!();
204        if start > 0 {
205            parts.push(
206                patch.wire_node(
207                    format!("{}.head", node.name),
208                    crate::ops::array::Slice::new(axis, 0, start),
209                    &[data_tap],
210                )?[0],
211            );
212        }
213        parts.push(patch.tap_model(model, node.inputs[2])?);
214        if start + len < data_shape[axis] {
215            parts.push(
216                patch.wire_node(
217                    format!("{}.tail", node.name),
218                    crate::ops::array::Slice::new(axis, start + len, data_shape[axis]),
219                    &[data_tap],
220                )?[0],
221            );
222        }
223        let wire = if parts.len() == 1 {
224            parts[0]
225        } else {
226            patch.wire_node(&node.name, crate::ops::array::TypedConcat::new(axis), &parts)?[0]
227        };
228        patch.shunt_outside(model, node.id.into(), wire)?;
229        Ok(Some(patch))
230    }
231}
232
233impl EvalOp for ScatterNd {
234    op_out_of_plan!();
235
236    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
237        let (data, indices, updates) = args_3!(inputs);
238        let indices = indices.cast_to::<i64>()?;
239        let indices = indices.to_plain_array_view::<i64>()?;
240        if data.datum_type() != updates.datum_type() {
241            bail!(
242                "Data and update must be of the same type, got {:?} and {:?}",
243                data.datum_type(),
244                updates.datum_type()
245            );
246        }
247        let mut data = data.into_tensor();
248        unsafe {
249            match self.reduction {
250                ScatterReduction::None => dispatch_datum_by_size!(
251                    Self::eval_t(data.datum_type())(&mut data, &indices, &updates)
252                )?,
253                reduction => dispatch_numbers!(Self::eval_t_reduce(data.datum_type())(
254                    &mut data, &indices, &updates, reduction
255                ))?,
256            }
257        }
258        Ok(tvec!(data.into_tvalue()))
259    }
260}
tract_core/ops/array/scatter_nd.rs

tract_core/ops/array/
scatter_nd.rs
