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

tract_core/ops/array/
dyn_slice.rs
