tract_core/ops/array/
slice.rs

1use crate::internal::*;
2use crate::num_traits::Zero;
3
4#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
5pub struct Slice {
6    pub axis: usize,
7    pub start: TDim,
8    pub end: TDim,
9}
10
11impl Slice {
12    pub fn new(axis: usize, start: impl ToDim, end: impl ToDim) -> Slice {
13        Slice { axis, start: start.to_dim(), end: end.to_dim() }
14    }
15
16    pub fn suffix(&self, name: &str) -> String {
17        format!("{}.axis{}_{}_{}", name, self.axis, self.start, self.end)
18    }
19
20    pub fn declutter_slice_after_slice(
21        &self,
22        model: &TypedModel,
23        node: &TypedNode,
24    ) -> TractResult<Option<TypedModelPatch>> {
25        let prec = model.node(node.inputs[0].node);
26        if let Some(other) = prec.op_as::<Slice>() {
27            if other.axis == self.axis {
28                return TypedModelPatch::replace_single_op(
29                    model,
30                    node,
31                    &prec.inputs,
32                    Slice {
33                        axis: self.axis,
34                        start: self.start.clone() + &other.start,
35                        end: self.end.clone() + &other.start,
36                    },
37                )
38                .map(Some);
39            }
40        }
41        Ok(None)
42    }
43}
44
45impl Op for Slice {
46    fn name(&self) -> StaticName {
47        "Slice".into()
48    }
49
50    fn info(&self) -> TractResult<Vec<String>> {
51        Ok(vec![format!("axis: {}, {}..{}", self.axis, self.start, self.end)])
52    }
53
54    op_as_typed_op!();
55
56    fn same_as(&self, other: &dyn Op) -> bool {
57        if let Some(other) = other.downcast_ref::<Self>() { other == self } else { false }
58    }
59}
60
61impl EvalOp for Slice {
62    fn is_stateless(&self) -> bool {
63        true
64    }
65
66    fn eval_with_session(
67        &self,
68        _node_id: usize,
69        session: &TurnState,
70        inputs: TVec<TValue>,
71    ) -> TractResult<TVec<TValue>> {
72        let input = args_1!(inputs);
73        let start = self.start.eval(&session.resolved_symbols).to_usize()?;
74        let end = self.end.eval(&session.resolved_symbols).to_usize()?;
75        eval_slice(&input, self.axis, start, end)
76    }
77}
78
79fn eval_slice(input: &Tensor, axis: usize, start: usize, end: usize) -> TractResult<TVec<TValue>> {
80    if end > input.shape()[axis] || start > end {
81        bail!("Invalid range {}..{} for slicing {:?} on axis {}", start, end, input, axis);
82    }
83    unsafe {
84        let mut shape: TVec<_> = input.shape().into();
85        shape[axis] = end - start;
86        let mut tensor = Tensor::uninitialized_dt(input.datum_type(), &shape)?;
87        tensor.assign_slice_unchecked(.., input, start..end, axis);
88        Ok(tvec!(tensor.into_tvalue()))
89    }
90}
91
92impl TypedOp for Slice {
93    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
94        anyhow::ensure!(inputs.len() == 1, "Slice has one single input");
95        if let (Ok(start), Ok(end), Ok(len)) =
96            (self.start.to_usize(), self.end.to_usize(), inputs[0].shape[self.axis].to_usize())
97        {
98            ensure!(start <= end);
99            ensure!(end <= len);
100        }
101        let mut fact = inputs[0].without_value();
102        fact.shape.set(self.axis, (self.end.clone() - &self.start).to_dim());
103        Ok(tvec!(fact))
104    }
105
106    fn axes_mapping(
107        &self,
108        inputs: &[&TypedFact],
109        outputs: &[&TypedFact],
110    ) -> TractResult<AxesMapping> {
111        let mut mapping = AxesMapping::disconnected(inputs, outputs)?;
112        for (axis, repr) in (0..inputs[0].rank()).zip('a'..) {
113            if self.axis != axis {
114                mapping = mapping
115                    .renaming((InOut::In(0), axis), repr)?
116                    .linking(repr, (InOut::Out(0), axis))?;
117            }
118        }
119        Ok(mapping)
120    }
121
122    fn change_axes(
123        &self,
124        model: &TypedModel,
125        node: &TypedNode,
126        _io: InOut,
127        change: &AxisOp,
128    ) -> TractResult<Option<AxisChangeConsequence>> {
129        if let Some(axis) = change.transform_axis(self.axis) {
130            if axis != self.axis {
131                Ok(Some(AxisChangeConsequence::new(
132                    model,
133                    node,
134                    Some(Box::new(Slice { axis, ..self.clone() }) as _),
135                    change,
136                )))
137            } else {
138                Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
139            }
140        } else {
141            Ok(None)
142        }
143    }
144
145    fn declutter(
146        &self,
147        model: &TypedModel,
148        node: &TypedNode,
149    ) -> TractResult<Option<TypedModelPatch>> {
150        if self.start.is_zero() && (self.end == model.outlet_fact(node.inputs[0])?.shape[self.axis])
151        {
152            TypedModelPatch::shunt_one_op(model, node)
153        } else if let Some(p) = self.declutter_slice_after_slice(model, node)? {
154            Ok(Some(p))
155        } else {
156            Ok(None)
157        }
158    }
159
160    fn concretize_dims(
161        &self,
162        _source: &TypedModel,
163        node: &TypedNode,
164        target: &mut TypedModel,
165        mapping: &HashMap<OutletId, OutletId>,
166        values: &SymbolValues,
167    ) -> TractResult<TVec<OutletId>> {
168        let op =
169            Slice { axis: self.axis, start: self.start.eval(values), end: self.end.eval(values) };
170        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
171        target.wire_node(&node.name, op, &inputs)
172    }
173
174    fn slice(
175        &self,
176        patch: &mut TypedModelPatch,
177        _model: &TypedModel,
178        node: &TypedNode,
179        _prefix: &str,
180        inputs: &[OutletId],
181        _output_axis: usize,
182        _start: &TDim,
183        _end: &TDim,
184    ) -> TractResult<Option<TVec<OutletId>>> {
185        patch.wire_node(&node.name, &node.op, inputs).map(Some)
186    }
187
188    as_op!();
189}
tract_core/ops/array/slice.rs

tract_core/ops/array/
slice.rs
