tract_hir/ops/array/
range.rs

1use tract_core::ops::cast::wire_cast;
2
3use crate::internal::*;
4
5#[derive(Debug, Default, Clone, new, Hash)]
6pub struct Range;
7
8impl Expansion for Range {
9    fn name(&self) -> StaticName {
10        "Range".into()
11    }
12
13    fn rules<'r, 'p: 'r, 's: 'r>(
14        &'s self,
15        s: &mut Solver<'r>,
16        inputs: &'p [TensorProxy],
17        outputs: &'p [TensorProxy],
18    ) -> InferenceResult {
19        check_input_arity(inputs, 3)?;
20        check_output_arity(outputs, 1)?;
21        s.given_3(
22            &inputs[0].datum_type,
23            &inputs[1].datum_type,
24            &inputs[2].datum_type,
25            move |s, dt0, dt1, dt2| {
26                let dt =
27                    DatumType::super_type_for([dt0, dt1, dt2]).context("No supertype found")?;
28                s.equals(dt, &outputs[0].datum_type)
29            },
30        )?;
31        s.equals(&inputs[0].rank, 0)?;
32        s.equals(&inputs[1].rank, 0)?;
33        s.equals(&inputs[2].rank, 0)?;
34        s.equals(&outputs[0].rank, 1)?;
35        s.given_3(&inputs[0].value, &inputs[1].value, &inputs[2].value, move |s, v0, v1, v2| {
36            let v0 = v0.cast_to::<TDim>()?;
37            let v1 = v1.cast_to::<TDim>()?;
38            let v2 = v2.cast_to::<i64>()?;
39            let out = (v1.to_scalar::<TDim>()?.clone() - v0.to_scalar::<TDim>()?)
40                .divceil(*v2.to_scalar::<i64>()? as _);
41            s.equals(&outputs[0].shape[0], out)
42        })?;
43        Ok(())
44    }
45
46    fn wire(
47        &self,
48        prefix: &str,
49        model: &mut TypedModel,
50        inputs: &[OutletId],
51    ) -> TractResult<TVec<OutletId>> {
52        let dt: DatumType = DatumType::super_type_for(
53            inputs.iter().map(|o| model.outlet_fact(*o).unwrap().datum_type),
54        )
55        .context("No supertype for inputs")?;
56        let inputs = wire_cast(prefix, model, inputs, dt)?;
57        let len = model.symbols.new_with_prefix("range");
58        model.wire_node(prefix, tract_core::ops::array::Range::new(len.into()), &inputs)
59    }
60}
tract_hir/ops/array/range.rs

tract_hir/ops/array/
range.rs
