tract_core/ops/array/
range.rs

1use crate::ops::cast::Cast;
2use tract_num_traits::AsPrimitive;
3use tract_num_traits::Zero;
4
5use crate::internal::*;
6
7use super::Slice;
8
9#[derive(Debug, Default, Clone, new, Hash, PartialEq, Eq)]
10pub struct Range {
11    len: TDim,
12}
13
14impl Op for Range {
15    fn name(&self) -> StaticName {
16        "Range".into()
17    }
18
19    op_as_typed_op!();
20}
21
22impl EvalOp for Range {
23    op_out_of_plan!();
24
25    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
26        let (start, end, step) = args_3!(inputs);
27        Ok(tvec!(self.make(&start, &end, &step, ctx.symbols)?.into_tvalue()))
28    }
29}
30
31impl Range {
32    fn make_t<T: Datum + for<'a> std::ops::Add<&'a T, Output = T>>(
33        start: &Tensor,
34        step: &Tensor,
35        len: usize,
36    ) -> TractResult<Tensor> {
37        unsafe {
38            let mut result = Tensor::uninitialized::<T>(&[len])?;
39            let mut v = start.try_as_plain()?.to_scalar::<T>()?.clone();
40            let step = step.try_as_plain()?.to_scalar::<T>()?;
41            {
42                let mut result_plain = result.try_as_plain_mut()?;
43                let slots = result_plain.as_slice_mut_unchecked::<T>().as_mut_ptr();
44                for i in 0..len {
45                    std::ptr::write(slots.add(i), v.clone());
46                    v = v + step;
47                }
48            }
49            Ok(result)
50        }
51    }
52
53    fn make(
54        &self,
55        start: &Tensor,
56        end: &Tensor,
57        step: &Tensor,
58        values: &SymbolValues,
59    ) -> TractResult<Tensor> {
60        if start.datum_type() == TDim::datum_type() {
61            let start = start.try_as_plain()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
62            let step = step.try_as_plain()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
63            let len = {
64                let end = end.try_as_plain()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
65                #[allow(clippy::cast_abs_to_unsigned)]
66                ((end - start).abs() as usize).divceil(step.abs() as usize)
67            };
68            Self::make_t::<i64>(&tensor0(start), &tensor0(step), len)
69        } else {
70            let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
71                self, start, end, step
72            ))?;
73            dispatch_numbers!(Self::make_t(start.datum_type())(start, step, len))
74        }
75    }
76
77    fn len_for_numbers<T: Datum + AsPrimitive<f64>>(
78        &self,
79        start: &Tensor,
80        end: &Tensor,
81        step: &Tensor,
82    ) -> TractResult<usize> {
83        let start = start.try_as_plain()?.to_scalar::<T>()?;
84        let end = end.try_as_plain()?.to_scalar::<T>()?;
85        let step = step.try_as_plain()?.to_scalar::<T>()?;
86        Ok(((end.as_() - start.as_()) / (step.as_())).ceil() as usize)
87    }
88}
89
90impl TypedOp for Range {
91    fn declutter(
92        &self,
93        model: &TypedModel,
94        node: &TypedNode,
95    ) -> TractResult<Option<TypedModelPatch>> {
96        rule_if_some!(succ = model.single_succ(node.id)?);
97        rule_if_some!(slice = succ.op_as::<Slice>());
98        rule_if!(slice.start.is_zero());
99        rule_if!(slice.end.is_zero());
100
101        let mut patch = TypedModelPatch::default();
102        let mut wire = patch.tap_model(model, node.inputs[0])?;
103        if model.outlet_fact(node.inputs[0])?.datum_type.is_tdim() {
104            wire = patch.wire_node(
105                format!("{}.cast-tdim", node.name),
106                Cast { to: DatumType::I64 },
107                &[wire],
108            )?[0];
109        }
110        let wire = patch.wire_node(&node.name, AxisOp::Add(0), &[wire])?;
111        patch.shunt_outside(model, succ.id.into(), wire[0])?;
112        Ok(Some(patch))
113    }
114
115    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
116        let [start, end, step] = inputs else {
117            bail!("Expects three inputs");
118        };
119        ensure!(start.datum_type() == end.datum_type());
120        ensure!(start.datum_type() == step.datum_type());
121        ensure!(start.shape.volume().is_one());
122        ensure!(end.shape.volume().is_one());
123        ensure!(step.shape.volume().is_one());
124        if let (Some(start), Some(end), Some(step)) = (&start.konst, &end.konst, &step.konst) {
125            if start.datum_type() == TDim::datum_type() {
126                let start_tdim = start.try_as_plain()?.to_scalar::<TDim>()?.clone();
127                let end_tdim = end.try_as_plain()?.to_scalar::<TDim>()?;
128                let step = step.cast_to_scalar::<i64>()?;
129                let len = if step < 0 {
130                    (start_tdim.clone() - end_tdim).divceil(-step as usize)
131                } else {
132                    (end_tdim.clone() - start_tdim.clone()).divceil(step as usize)
133                };
134                let mut fact = DatumType::I64.fact([len]);
135                if let Some(scope) = start_tdim.find_scope().or_else(|| end_tdim.find_scope()) {
136                    let x0 = TDim::Sym(scope.coord_sym(0));
137                    let term = if step == 1 { x0 } else { TDim::MulInt(step, Box::new(x0)) };
138                    fact.uniform_tdim = Some((start_tdim + term).reduce());
139                }
140                Ok(tvec!(fact))
141            } else {
142                let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
143                    self, start, end, step
144                ))?
145                .to_dim();
146                Ok(tvec!(start.datum_type().fact([len])))
147            }
148        } else {
149            let mut fact = start.datum_type.fact(std::slice::from_ref(&self.len));
150            if let (Some(s), Some(k)) = (&start.uniform_tdim, &step.uniform_tdim)
151                && let Some(scope) = self.len.find_scope()
152            {
153                let x0 = TDim::Sym(scope.coord_sym(0));
154                let term = match k {
155                    TDim::Val(1) => x0,
156                    TDim::Val(v) => TDim::MulInt(*v, Box::new(x0)),
157                    other => TDim::Mul(vec![other.clone(), x0]),
158                };
159                fact.uniform_tdim = Some((s.clone() + term).reduce());
160            }
161            Ok(tvec!(fact))
162        }
163    }
164
165    as_op!();
166}
tract_core/ops/array/range.rs

tract_core/ops/array/
range.rs
