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 && other.axis == self.axis
28 {
29 return TypedModelPatch::replace_single_op(
30 model,
31 node,
32 &prec.inputs,
33 Slice {
34 axis: self.axis,
35 start: self.start.clone() + &other.start,
36 end: self.end.clone() + &other.start,
37 },
38 )
39 .map(Some);
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
57impl EvalOp for Slice {
58 op_out_of_plan!();
59
60 fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
61 let input = args_1!(inputs);
62 let start = self.start.eval(ctx.symbols).to_usize()?;
63 let end = self.end.eval(ctx.symbols).to_usize()?;
64 eval_slice(&input, self.axis, start, end)
65 }
66}
67
68fn eval_slice(input: &Tensor, axis: usize, start: usize, end: usize) -> TractResult<TVec<TValue>> {
69 if end > input.shape()[axis] || start > end {
70 bail!("Invalid range {}..{} for slicing {:?} on axis {}", start, end, input, axis);
71 }
72 unsafe {
73 let mut shape: TVec<_> = input.shape().into();
74 shape[axis] = end - start;
75 let mut tensor = Tensor::uninitialized_dt(input.datum_type(), &shape)?;
76 tensor.assign_slice_unchecked(.., input, start..end, axis);
77 Ok(tvec!(tensor.into_tvalue()))
78 }
79}
80
81impl TypedOp for Slice {
82 fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
83 anyhow::ensure!(inputs.len() == 1, "Slice has one single input");
84 if let (Some(start), Some(end), Some(len)) =
85 (self.start.as_usize(), self.end.as_usize(), inputs[0].shape[self.axis].as_usize())
86 {
87 ensure!(start <= end);
88 ensure!(end <= len);
89 }
90 let mut fact = inputs[0].without_value();
91 fact.shape.set(self.axis, (self.end.clone() - &self.start).to_dim());
92 Ok(tvec!(fact))
93 }
94
95 fn input_roi(
96 &self,
97 model: &TypedModel,
98 node: &TypedNode,
99 ) -> TractResult<Option<TVec<Option<TDim>>>> {
100 let output_fact = model.outlet_fact(OutletId::new(node.id, 0))?;
101 rule_if_some!(roi = &output_fact.region_of_interest);
102 if self.start.is_zero() {
103 return Ok(Some(tvec![Some(roi.clone())]));
104 }
105 if let Some(sym) = roi
107 .symbols()
108 .into_iter()
109 .find(|s| crate::ops::logic::sym_to_coord_axis(s) == Some(self.axis))
110 {
111 let shifted = TDim::Sym(sym.clone()) + self.start.clone();
112 if let Ok(input_roi) = roi.substitute(&sym, &shifted) {
113 return Ok(Some(tvec![Some(input_roi)]));
114 }
115 }
116 Ok(Some(tvec![Some(roi.clone())]))
118 }
119
120 fn axes_mapping(
121 &self,
122 inputs: &[&TypedFact],
123 outputs: &[&TypedFact],
124 ) -> TractResult<AxesMapping> {
125 let mut mapping = AxesMapping::disconnected(inputs, outputs)?;
126 for (axis, repr) in (0..inputs[0].rank()).zip('a'..) {
127 if self.axis != axis {
128 mapping = mapping
129 .renaming((InOut::In(0), axis), repr)?
130 .linking(repr, (InOut::Out(0), axis))?;
131 }
132 }
133 Ok(mapping)
134 }
135
136 fn change_axes(
137 &self,
138 model: &TypedModel,
139 node: &TypedNode,
140 _io: InOut,
141 change: &AxisOp,
142 ) -> TractResult<Option<AxisChangeConsequence>> {
143 if let Some(axis) = change.transform_axis(self.axis) {
144 if axis != self.axis {
145 Ok(Some(AxisChangeConsequence::new(
146 model,
147 node,
148 Some(Box::new(Slice { axis, ..self.clone() }) as _),
149 change,
150 )))
151 } else {
152 Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
153 }
154 } else {
155 Ok(None)
156 }
157 }
158
159 fn declutter(
160 &self,
161 model: &TypedModel,
162 node: &TypedNode,
163 ) -> TractResult<Option<TypedModelPatch>> {
164 if self.start.is_zero() && (self.end == model.outlet_fact(node.inputs[0])?.shape[self.axis])
165 {
166 TypedModelPatch::shunt_one_op(model, node)
167 } else if let Some(p) = self.declutter_slice_after_slice(model, node)? {
168 Ok(Some(p))
169 } else {
170 Ok(None)
171 }
172 }
173
174 fn set_symbols(
175 &self,
176 _source: &TypedModel,
177 node: &TypedNode,
178 target: &mut TypedModel,
179 mapping: &HashMap<OutletId, OutletId>,
180 subs: &HashMap<Symbol, TDim>,
181 ) -> TractResult<TVec<OutletId>> {
182 let op = Slice {
183 axis: self.axis,
184 start: self.start.substitute_all(subs)?,
185 end: self.end.substitute_all(subs)?,
186 };
187 let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
188 target.wire_node(&node.name, op, &inputs)
189 }
190
191 fn slice(
192 &self,
193 patch: &mut TypedModelPatch,
194 _model: &TypedModel,
195 node: &TypedNode,
196 _prefix: &str,
197 inputs: &[OutletId],
198 _output_axis: usize,
199 _start: &TDim,
200 _end: &TDim,
201 ) -> TractResult<Option<TVec<OutletId>>> {
202 patch.wire_node(&node.name, &node.op, inputs).map(Some)
203 }
204
205 as_op!();
206}