tract_core/ops/array/
cumsum.rs

1use super::MultiBroadcastTo;
2use crate::internal::*;
3use tract_itertools::Itertools;
4use tract_num_traits::Zero;
5
6/// Cumulative sum along `axis`.
7///
8/// Output `i` is the sum of inputs `0..=i` (`0..i` when `exclusive`), counted from the
9/// end of the axis when `reverse`. Shape and type are the input's; every numeric type is
10/// accepted. The axes mapping leaves `axis` disconnected so no rule slices, pulses or
11/// moves through it as if positions were independent. Declutter folds a sum over a
12/// length-1 axis to the identity, or to zeros when `exclusive`.
13#[derive(Clone, Debug, Hash, PartialEq, Eq)]
14pub struct CumSum {
15    pub axis: usize,
16    pub exclusive: bool,
17    pub reverse: bool,
18}
19
20impl CumSum {
21    fn eval_t<T: Datum + Copy + std::ops::AddAssign + Zero>(
22        &self,
23        tensor: &mut Tensor,
24    ) -> TractResult<()> {
25        let shape = tensor.shape().to_vec();
26        let len = shape[self.axis];
27        let inner: usize = shape[self.axis + 1..].iter().product();
28        if len == 0 || inner == 0 {
29            return Ok(());
30        }
31        let mut plain = tensor.try_as_plain_ram_mut()?;
32        let data = plain.as_slice_mut::<T>()?;
33        for block in data.chunks_exact_mut(len * inner) {
34            if self.reverse {
35                for i in (0..len - 1).rev() {
36                    let (head, tail) = block.split_at_mut((i + 1) * inner);
37                    add_row(&mut head[i * inner..], &tail[..inner]);
38                }
39                if self.exclusive {
40                    block.copy_within(inner.., 0);
41                    block[(len - 1) * inner..].fill(T::zero());
42                }
43            } else {
44                for i in 1..len {
45                    let (head, tail) = block.split_at_mut(i * inner);
46                    add_row(&mut tail[..inner], &head[(i - 1) * inner..]);
47                }
48                if self.exclusive {
49                    block.copy_within(..(len - 1) * inner, inner);
50                    block[..inner].fill(T::zero());
51                }
52            }
53        }
54        Ok(())
55    }
56}
57
58fn add_row<T: Copy + std::ops::AddAssign>(acc: &mut [T], row: &[T]) {
59    acc.iter_mut().zip(row).for_each(|(a, r)| *a += *r);
60}
61
62impl Op for CumSum {
63    fn name(&self) -> StaticName {
64        "CumSum".into()
65    }
66
67    fn info(&self) -> TractResult<Vec<String>> {
68        Ok(vec![format!(
69            "axis: {} exclusive: {} reverse: {}",
70            self.axis, self.exclusive, self.reverse
71        )])
72    }
73
74    op_as_typed_op!();
75}
76
77impl EvalOp for CumSum {
78    op_out_of_plan!();
79
80    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
81        let mut tensor = args_1!(inputs).into_tensor();
82        dispatch_numbers!(Self::eval_t(tensor.datum_type())(self, &mut tensor))?;
83        Ok(tvec!(tensor.into_tvalue()))
84    }
85}
86
87impl TypedOp for CumSum {
88    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
89        ensure!(self.axis < inputs[0].rank());
90        ensure!(!inputs[0].datum_type.is_quantized(), "CumSum does not take quantized inputs");
91        Ok(tvec!(inputs[0].datum_type.fact(inputs[0].shape.clone())))
92    }
93
94    fn axes_mapping(
95        &self,
96        inputs: &[&TypedFact],
97        outputs: &[&TypedFact],
98    ) -> TractResult<AxesMapping> {
99        let mut letters = 'a'..;
100        let axes = (0..inputs[0].rank())
101            .flat_map(|ix| {
102                if ix == self.axis {
103                    tvec!(
104                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
105                            .input(0, ix),
106                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
107                            .output(0, ix),
108                    )
109                } else {
110                    tvec!(
111                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
112                            .input(0, ix)
113                            .output(0, ix)
114                    )
115                }
116            })
117            .collect_vec();
118        AxesMapping::new(1, 1, axes)
119    }
120
121    fn declutter(
122        &self,
123        model: &TypedModel,
124        node: &TypedNode,
125    ) -> TractResult<Option<TypedModelPatch>> {
126        let fact = model.outlet_fact(node.inputs[0])?;
127        rule_if!(fact.shape[self.axis].is_one());
128        if !self.exclusive {
129            return TypedModelPatch::shunt_one_op(model, node);
130        }
131        // An exclusive sum over a single position is all zeros.
132        let zero = Tensor::zero_dt(fact.datum_type, &tvec![1; fact.rank()])?;
133        let mut patch = TypedModelPatch::default();
134        let zero = patch.add_const(format!("{}.zero", node.name), zero)?;
135        let wire =
136            patch.wire_node(&node.name, MultiBroadcastTo::new(fact.shape.clone()), &[zero])?;
137        patch.shunt_outside(model, node.id.into(), wire[0])?;
138        Ok(Some(patch))
139    }
140
141    fn change_axes(
142        &self,
143        model: &TypedModel,
144        node: &TypedNode,
145        _io: InOut,
146        change: &AxisOp,
147    ) -> TractResult<Option<AxisChangeConsequence>> {
148        rule_if_some!(axis = change.transform_axis(self.axis));
149        let op = Some(Box::new(Self { axis, ..self.clone() }) as _);
150        Ok(Some(AxisChangeConsequence::new(model, node, op, change)))
151    }
152
153    fn slice(
154        &self,
155        patch: &mut TypedModelPatch,
156        _model: &TypedModel,
157        node: &TypedNode,
158        _prefix: &str,
159        inputs: &[OutletId],
160        output_axis: usize,
161        _start: &TDim,
162        _end: &TDim,
163    ) -> TractResult<Option<TVec<OutletId>>> {
164        rule_if!(output_axis != self.axis);
165        patch.wire_node(&node.name, &node.op, inputs).map(Some)
166    }
167
168    as_op!();
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    fn declutter_single_step(exclusive: bool) -> TractResult<TypedModel> {
176        let mut model = TypedModel::default();
177        let t = model.symbols.sym("T");
178        let input = model.add_source("input", f32::fact(&[1.to_dim(), t.to_dim()]))?;
179        let op = CumSum { axis: 0, exclusive, reverse: false };
180        let output = model.wire_node("cumsum", op, &[input])?;
181        model.select_output_outlets(&output)?;
182        model.into_decluttered()
183    }
184
185    #[test]
186    fn single_step_inclusive_is_identity() -> TractResult<()> {
187        let model = declutter_single_step(false)?;
188        assert_eq!(model.nodes().len(), 1);
189        Ok(())
190    }
191
192    #[test]
193    fn single_step_exclusive_is_zeros() -> TractResult<()> {
194        let model = declutter_single_step(true)?;
195        assert!(model.nodes().iter().all(|n| !n.op_is::<CumSum>()));
196        let input = tensor2(&[[1f32, 2., 3.]]);
197        let output = model.into_runnable()?.run(tvec!(input.into_tvalue()))?;
198        output[0].close_enough(&tensor2(&[[0f32, 0., 0.]]), Approximation::Exact)?;
199        Ok(())
200    }
201}
tract_core/ops/array/cumsum.rs

tract_core/ops/array/
cumsum.rs
