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 mints_runtime_symbols(&self) -> bool {
16        true
17    }
18
19    fn name(&self) -> StaticName {
20        "Range".into()
21    }
22
23    op_as_typed_op!();
24}
25
26impl EvalOp for Range {
27    op_out_of_plan!();
28
29    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
30        let (start, end, step) = args_3!(inputs);
31        Ok(tvec!(self.make(&start, &end, &step, ctx.symbols)?.into_tvalue()))
32    }
33}
34
35impl Range {
36    fn make_t<T: Datum + for<'a> std::ops::Add<&'a T, Output = T>>(
37        start: &Tensor,
38        step: &Tensor,
39        len: usize,
40    ) -> TractResult<Tensor> {
41        unsafe {
42            let mut result = Tensor::uninitialized::<T>(&[len])?;
43            let mut v = start.try_as_plain_ram()?.to_scalar::<T>()?.clone();
44            let step = step.try_as_plain_ram()?.to_scalar::<T>()?;
45            {
46                let mut result_plain = result.try_as_plain_ram_mut()?;
47                let slots = result_plain.as_slice_mut_unchecked::<T>().as_mut_ptr();
48                for i in 0..len {
49                    std::ptr::write(slots.add(i), v.clone());
50                    v = v + step;
51                }
52            }
53            Ok(result)
54        }
55    }
56
57    fn make(
58        &self,
59        start: &Tensor,
60        end: &Tensor,
61        step: &Tensor,
62        values: &SymbolValues,
63    ) -> TractResult<Tensor> {
64        if start.datum_type() == TDim::datum_type() {
65            // TDim inputs always yield an i64 tensor: both branches of
66            // output_facts declare i64 for TDim inputs (the symbolic value
67            // rides in uniform_tdim), and the plan's fact-vs-value assertions
68            // check that.
69            let start = start.try_as_plain_ram()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
70            let step = step.try_as_plain_ram()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
71            let end = end.try_as_plain_ram()?.to_scalar::<TDim>()?.eval(values).to_i64()?;
72            let len = {
73                // i128 intermediate: end - start can overflow i64
74                let span = usize::try_from((end as i128 - start as i128).unsigned_abs()).ok();
75                let step = usize::try_from(step.unsigned_abs()).ok();
76                match (span, step) {
77                    (Some(span), Some(step)) if step > 0 => span.divceil(step),
78                    _ => bail!("Range span or step out of usize bounds"),
79                }
80            };
81            Self::make_t::<i64>(&tensor0(start), &tensor0(step), len)
82        } else {
83            let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
84                self, start, end, step
85            ))?;
86            dispatch_numbers!(Self::make_t(start.datum_type())(start, step, len))
87        }
88    }
89
90    fn len_for_numbers<T: Datum + AsPrimitive<f64>>(
91        &self,
92        start: &Tensor,
93        end: &Tensor,
94        step: &Tensor,
95    ) -> TractResult<usize> {
96        let start = start.try_as_plain_ram()?.to_scalar::<T>()?;
97        let end = end.try_as_plain_ram()?.to_scalar::<T>()?;
98        let step = step.try_as_plain_ram()?.to_scalar::<T>()?;
99        Ok(((end.as_() - start.as_()) / (step.as_())).ceil() as usize)
100    }
101}
102
103impl TypedOp for Range {
104    fn declutter(
105        &self,
106        model: &TypedModel,
107        node: &TypedNode,
108    ) -> TractResult<Option<TypedModelPatch>> {
109        rule_if_some!(succ = model.single_succ(node.id)?);
110        rule_if_some!(slice = succ.op_as::<Slice>());
111        rule_if!(slice.start.is_zero());
112        rule_if!(slice.end.is_zero());
113
114        let mut patch = TypedModelPatch::default();
115        let mut wire = patch.tap_model(model, node.inputs[0])?;
116        if model.outlet_fact(node.inputs[0])?.datum_type.is_tdim() {
117            wire = patch.wire_node(
118                format!("{}.cast-tdim", node.name),
119                Cast { to: DatumType::I64 },
120                &[wire],
121            )?[0];
122        }
123        let wire = patch.wire_node(&node.name, AxisOp::Add(0), &[wire])?;
124        patch.shunt_outside(model, succ.id.into(), wire[0])?;
125        Ok(Some(patch))
126    }
127
128    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
129        let [start, end, step] = inputs else {
130            bail!("Expects three inputs");
131        };
132        ensure!(start.datum_type() == end.datum_type());
133        ensure!(start.datum_type() == step.datum_type());
134        ensure!(start.shape.volume().is_one());
135        ensure!(end.shape.volume().is_one());
136        ensure!(step.shape.volume().is_one());
137        if let (Some(start), Some(end), Some(step)) = (&start.konst, &end.konst, &step.konst) {
138            if start.datum_type() == TDim::datum_type() {
139                let start_tdim = start.try_as_plain_ram()?.to_scalar::<TDim>()?.clone();
140                let end_tdim = end.try_as_plain_ram()?.to_scalar::<TDim>()?;
141                let step = step.cast_to_scalar::<i64>()?;
142                let len = if step < 0 {
143                    (start_tdim.clone() - end_tdim).divceil(-step as usize)
144                } else {
145                    (end_tdim.clone() - start_tdim.clone()).divceil(step as usize)
146                };
147                let mut fact = DatumType::I64.fact([len]);
148                if let Some(scope) = start_tdim.find_scope().or_else(|| end_tdim.find_scope()) {
149                    let x0 = TDim::Sym(scope.coord_sym(0));
150                    let term = if step == 1 { x0 } else { TDim::MulInt(step, Box::new(x0)) };
151                    fact.uniform_tdim = Some((start_tdim + term).reduce());
152                }
153                Ok(tvec!(fact))
154            } else {
155                let len = dispatch_numbers!(Self::len_for_numbers(start.datum_type())(
156                    self, start, end, step
157                ))?
158                .to_dim();
159                Ok(tvec!(start.datum_type().fact([len])))
160            }
161        } else {
162            // TDim inputs yield i64, like the konst branch above: the
163            // symbolic per-element value rides in uniform_tdim, and
164            // Range::make materializes i64 regardless of the input datum
165            // type (a contract pulse-opl relies on).
166            let dt = if start.datum_type.is_tdim() { i64::datum_type() } else { start.datum_type };
167            let mut fact = dt.fact(std::slice::from_ref(&self.len));
168            if let (Some(s), Some(k)) = (&start.uniform_tdim, &step.uniform_tdim)
169                && let Some(scope) = self.len.find_scope()
170            {
171                let x0 = TDim::Sym(scope.coord_sym(0));
172                let term = match k {
173                    TDim::Val(1) => x0,
174                    TDim::Val(v) => TDim::MulInt(*v, Box::new(x0)),
175                    other => TDim::Mul(vec![other.clone(), x0]),
176                };
177                fact.uniform_tdim = Some((s.clone() + term).reduce());
178            }
179            Ok(tvec!(fact))
180        }
181    }
182
183    as_op!();
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    fn tdims_model(model: &mut TypedModel, end: OutletId) -> TractResult<TVec<OutletId>> {
191        let start = model.wire_node(
192            "start",
193            crate::ops::konst::Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?,
194            &[],
195        )?;
196        let step = model.wire_node(
197            "step",
198            crate::ops::konst::Const::new(tensor0(TDim::Val(1)).into_arc_tensor())?,
199            &[],
200        )?;
201        let s = model.symbols.sym("S");
202        model.wire_node("range", Range::new(s.to_dim()), &[start[0], end, step[0]])
203    }
204
205    // sdpa-style: all-konst TDim scalars (the konst branch of output_facts).
206    // Both the fact and the evaluated tensor must be i64. S is bound at
207    // runtime through the definer source' shape.
208    #[test]
209    fn konst_tdim_inputs_yield_i64() -> TractResult<()> {
210        let mut model = TypedModel::default();
211        let s = model.symbols.sym("S");
212        let definer = model.add_source("definer", f32::datum_type().fact([s.to_dim()]))?;
213        let end = model.wire_node(
214            "end",
215            crate::ops::konst::Const::new(tensor0(TDim::Sym(s)).into_arc_tensor())?,
216            &[],
217        )?;
218        let range = tdims_model(&mut model, end[0])?;
219        model.select_output_outlets(&[range[0], definer])?;
220        assert_eq!(model.outlet_fact(range[0])?.datum_type, i64::datum_type());
221        let found = crate::internal::TypedSimplePlan::new(model)?
222            .run(tvec!(tensor1(&[0f32; 5]).into_tvalue()))?;
223        assert_eq!(*found[0], tensor1(&[0i64, 1, 2, 3, 4]));
224        Ok(())
225    }
226
227    // dynamic end (the else branch of output_facts): i64 fact, i64 values,
228    // symbolic per-element value still tracked as uniform_tdim.
229    #[test]
230    fn dynamic_tdim_end_yields_i64() -> TractResult<()> {
231        let mut model = TypedModel::default();
232        let end = model.add_source("T_dyn", TDim::datum_type().scalar_fact())?;
233        let range = tdims_model(&mut model, end)?;
234        model.select_output_outlets(&range)?;
235        let fact = model.outlet_fact(range[0])?;
236        assert_eq!(fact.datum_type, i64::datum_type());
237        assert!(fact.uniform_tdim.is_some());
238        let found = crate::internal::TypedSimplePlan::new(model)?
239            .run(tvec!(tensor0(TDim::Val(5)).into_tvalue()))?;
240        assert_eq!(*found[0], tensor1(&[0i64, 1, 2, 3, 4]));
241        Ok(())
242    }
243}
tract_core/ops/array/range.rs

tract_core/ops/array/
range.rs
