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) -> Cow<str> {
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>() {
58            other == self
59        } else {
60            false
61        }
62    }
63}
64
65impl EvalOp for Slice {
66    fn is_stateless(&self) -> bool {
67        true
68    }
69
70    fn eval_with_session(
71        &self,
72        session: &SessionState,
73        inputs: TVec<TValue>,
74    ) -> TractResult<TVec<TValue>> {
75        let input = args_1!(inputs);
76        let start = self.start.eval(&session.resolved_symbols).to_usize()?;
77        let end = self.end.eval(&session.resolved_symbols).to_usize()?;
78        eval_slice(&input, self.axis, start, end)
79    }
80}
81
82fn eval_slice(input: &Tensor, axis: usize, start: usize, end: usize) -> TractResult<TVec<TValue>> {
83    if end > input.shape()[axis] || start > end {
84        bail!("Invalid range {}..{} for slicing {:?} on axis {}", start, end, input, axis);
85    }
86    unsafe {
87        let mut shape: TVec<_> = input.shape().into();
88        shape[axis] = end - start;
89        let mut tensor = Tensor::uninitialized_dt(input.datum_type(), &shape)?;
90        tensor.assign_slice_unchecked(.., input, start..end, axis);
91        Ok(tvec!(tensor.into_tvalue()))
92    }
93}
94
95impl TypedOp for Slice {
96    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
97        anyhow::ensure!(inputs.len() == 1, "Slice has one single input");
98        if let (Ok(start), Ok(end), Ok(len)) =
99            (self.start.to_usize(), self.end.to_usize(), inputs[0].shape[self.axis].to_usize())
100        {
101            ensure!(start <= end);
102            ensure!(end <= len);
103        }
104        let mut fact = inputs[0].without_value();
105        fact.shape.set(self.axis, (self.end.clone() - &self.start).to_dim());
106        Ok(tvec!(fact))
107    }
108
109    fn axes_mapping(
110        &self,
111        inputs: &[&TypedFact],
112        outputs: &[&TypedFact],
113    ) -> TractResult<AxesMapping> {
114        let mut mapping = AxesMapping::disconnected(inputs, outputs)?;
115        for (axis, repr) in (0..inputs[0].rank()).zip('a'..) {
116            if self.axis != axis {
117                mapping = mapping
118                    .renaming((InOut::In(0), axis), repr)?
119                    .linking(repr, (InOut::Out(0), axis))?;
120            }
121        }
122        Ok(mapping)
123    }
124
125    fn change_axes(
126        &self,
127        model: &TypedModel,
128        node: &TypedNode,
129        _io: InOut,
130        change: &AxisOp,
131    ) -> TractResult<Option<AxisChangeConsequence>> {
132        if let Some(axis) = change.transform_axis(self.axis) {
133            if axis != self.axis {
134                Ok(Some(AxisChangeConsequence::new(
135                    model,
136                    node,
137                    Some(Box::new(Slice { axis, ..self.clone() }) as _),
138                    change,
139                )))
140            } else {
141                Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
142            }
143        } else {
144            Ok(None)
145        }
146    }
147
148    fn declutter(
149        &self,
150        model: &TypedModel,
151        node: &TypedNode,
152    ) -> TractResult<Option<TypedModelPatch>> {
153        if self.start.is_zero() && (self.end == model.outlet_fact(node.inputs[0])?.shape[self.axis])
154        {
155            TypedModelPatch::shunt_one_op(model, node)
156        } else if let Some(p) = self.declutter_slice_after_slice(model, node)? {
157            Ok(Some(p))
158        } else {
159            Ok(None)
160        }
161    }
162
163    fn concretize_dims(
164        &self,
165        _source: &TypedModel,
166        node: &TypedNode,
167        target: &mut TypedModel,
168        mapping: &HashMap<OutletId, OutletId>,
169        values: &SymbolValues,
170    ) -> TractResult<TVec<OutletId>> {
171        let op =
172            Slice { axis: self.axis, start: self.start.eval(values), end: self.end.eval(values) };
173        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
174        target.wire_node(&node.name, op, &inputs)
175    }
176
177    as_op!();
178}
tract_core/ops/array/slice.rs

tract_core/ops/array/
slice.rs
