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
115impl TypedOp for ScatterNd {
116    as_op!();
117
118    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
119        Ok(tvec!(inputs[0].datum_type.fact(inputs[0].shape.to_tvec())))
120    }
121}
122
123impl EvalOp for ScatterNd {
124    fn is_stateless(&self) -> bool {
125        true
126    }
127
128    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
129        let (data, indices, updates) = args_3!(inputs);
130        let indices = indices.cast_to::<i64>()?;
131        let indices = indices.to_plain_array_view::<i64>()?;
132        if data.datum_type() != updates.datum_type() {
133            bail!(
134                "Data and update must be of the same type, got {:?} and {:?}",
135                data.datum_type(),
136                updates.datum_type()
137            );
138        }
139        let mut data = data.into_tensor();
140        unsafe {
141            match self.reduction {
142                ScatterReduction::None => dispatch_datum_by_size!(
143                    Self::eval_t(data.datum_type())(&mut data, &indices, &updates)
144                )?,
145                reduction => dispatch_numbers!(Self::eval_t_reduce(data.datum_type())(
146                    &mut data, &indices, &updates, reduction
147                ))?,
148            }
149        }
150        Ok(tvec!(data.into_tvalue()))
151    }
152}
tract_core/ops/array/scatter_nd.rs

tract_core/ops/array/
scatter_nd.rs
