tract_core/ops/array/
dyn_slice.rs

1use crate::internal::*;
2
3#[derive(Debug, Clone, PartialEq, Eq, Hash, new)]
4pub struct DynSlice {
5    pub axis: usize,
6    pub len: TDim,
7}
8
9impl DynSlice {
10    pub fn suffix(&self) -> String {
11        format!("axis{}", self.axis)
12    }
13}
14
15impl Op for DynSlice {
16    fn mints_runtime_symbols(&self) -> bool {
17        true
18    }
19
20    fn name(&self) -> StaticName {
21        "DynSlice".into()
22    }
23
24    fn info(&self) -> TractResult<Vec<String>> {
25        Ok(vec![format!("axis: {}", self.axis)])
26    }
27
28    op_as_typed_op!();
29}
30
31impl EvalOp for DynSlice {
32    op_out_of_plan!();
33
34    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
35        let start = inputs[1]
36            .cast_to::<TDim>()?
37            .try_as_plain_ram()?
38            .to_scalar::<TDim>()?
39            .eval(ctx.symbols)
40            .to_usize()?;
41        let end = inputs[2]
42            .cast_to::<TDim>()?
43            .try_as_plain_ram()?
44            .to_scalar::<TDim>()?
45            .eval(ctx.symbols)
46            .to_usize()?;
47        ensure!(start <= end);
48        if let Ok(len) = self.len.eval(ctx.symbols).to_usize() {
49            ensure!(start + len == end);
50        }
51        let slice = inputs[0].slice(self.axis, start, end)?;
52        Ok(tvec!(slice.into()))
53    }
54}
55
56impl TypedOp for DynSlice {
57    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
58        ensure!(inputs.len() == 3);
59        let mut fact = inputs[0].without_value();
60        fact.shape.set(self.axis, self.len.clone());
61        Ok(tvec!(fact))
62    }
63
64    fn axes_mapping(
65        &self,
66        inputs: &[&TypedFact],
67        _outputs: &[&TypedFact],
68    ) -> TractResult<AxesMapping> {
69        AxesMapping::natural_for_rank(1, 1, inputs[0].rank())?
70            .with_extra_input(1)?
71            .with_extra_input(2)
72    }
73
74    fn change_axes(
75        &self,
76        model: &TypedModel,
77        node: &TypedNode,
78        io: InOut,
79        change: &AxisOp,
80    ) -> TractResult<Option<AxisChangeConsequence>> {
81        rule_if!(io != InOut::In(1) && io != InOut::In(2));
82        rule_if_some!(axis = change.transform_axis(self.axis));
83        if axis != self.axis {
84            Ok(Some(AxisChangeConsequence::new(
85                model,
86                node,
87                Some(Box::new(DynSlice { axis, ..self.clone() }) as _),
88                change,
89            )))
90        } else {
91            Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
92        }
93    }
94
95    fn declutter(
96        &self,
97        model: &TypedModel,
98        node: &TypedNode,
99    ) -> TractResult<Option<TypedModelPatch>> {
100        let inputs = model.node_input_facts(node.id)?;
101        rule_if_some!(start = &inputs[1].konst);
102        rule_if_some!(end = &inputs[2].konst);
103        let start = start.cast_to::<TDim>()?.try_as_plain_ram()?.to_scalar::<TDim>()?.clone();
104        let end = end.cast_to::<TDim>()?.try_as_plain_ram()?.to_scalar::<TDim>()?.clone();
105
106        Ok(Some(TypedModelPatch::replace_single_op(
107            model,
108            node,
109            &[node.inputs[0]],
110            crate::ops::array::Slice { axis: self.axis, start, end },
111        )?))
112    }
113
114    fn set_symbols(
115        &self,
116        _source: &TypedModel,
117        node: &TypedNode,
118        target: &mut TypedModel,
119        mapping: &HashMap<OutletId, OutletId>,
120        subs: &HashMap<Symbol, TDim>,
121    ) -> TractResult<TVec<OutletId>> {
122        let op = DynSlice { axis: self.axis, len: self.len.substitute_all(subs)? };
123        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
124        target.wire_node(&node.name, op, &inputs)
125    }
126
127    as_op!();
128}
tract_core/ops/array/dyn_slice.rs

tract_core/ops/array/
dyn_slice.rs
