1use crate::internal::*;
2
3#[derive(Debug, Clone, PartialEq, Eq, Hash)]
4pub enum PadMode {
5 Constant(Arc<Tensor>),
6 Reflect,
7 Edge,
8}
9
10impl Default for PadMode {
11 fn default() -> PadMode {
12 PadMode::Constant(Arc::new(0.0f32.into()))
13 }
14}
15
16#[derive(Debug, Clone, new, Default, Hash, PartialEq, Eq)]
17pub struct Pad {
18 pub pads: Vec<(usize, usize)>,
19 pub mode: PadMode,
20}
21
22impl Pad {
23 fn eval_t<T>(&self, input_tensor: TValue) -> TractResult<TValue>
24 where
25 T: Copy + Datum,
26 {
27 use tract_ndarray::*;
28 let input = input_tensor.to_plain_array_view::<T>()?;
29 let output_shape: Vec<usize> =
30 input.shape().iter().zip(self.pads.iter()).map(|(&d, &(a, b))| d + a + b).collect();
31 let element = match &self.mode {
32 PadMode::Constant(f) => f.cast_to_scalar::<T>()?,
33 _ => T::default(),
34 };
35 let mut output = ArrayD::<T>::from_elem(output_shape, element);
36 let slice_spec: Vec<SliceInfoElem> = self
37 .pads
38 .iter()
39 .map(|&(a, b)| SliceInfoElem::Slice {
40 start: a as isize,
41 end: if b != 0 { Some(-(b as isize)) } else { None },
42 step: 1,
43 })
44 .collect();
45 let slice_info = SliceInfo::<_, IxDyn, IxDyn>::try_from(slice_spec).unwrap();
46 output.slice_mut(slice_info.as_ref()).assign(&input);
47 if self.mode == PadMode::Reflect || self.mode == PadMode::Edge {
48 for (ax, &(bef, aft)) in self.pads.iter().enumerate() {
49 let axis = Axis(ax);
50 let dim = output.shape()[ax];
51 {
52 let (mut pad, data) = output.view_mut().split_at(axis, bef);
53 for i in 0..bef {
54 let mut target = pad.slice_axis_mut(axis, Slice::from(i..i + 1));
55 let source_slice = match self.mode {
56 PadMode::Edge => 0,
57 PadMode::Reflect => bef - i,
58 _ => panic!(),
59 };
60 let source =
61 data.slice_axis(axis, Slice::from(source_slice..source_slice + 1));
62 target.assign(&source);
63 }
64 }
65 {
66 let (data, mut pad) = output.view_mut().split_at(axis, dim - aft);
67 for i in 0..aft {
68 let mut target = pad.slice_axis_mut(axis, Slice::from(i..i + 1));
69 let source_slice = match self.mode {
70 PadMode::Edge => dim - aft - 1,
71 PadMode::Reflect => dim - aft - 2 - i,
72 _ => panic!(),
73 };
74 let source =
75 data.slice_axis(axis, Slice::from(source_slice..source_slice + 1));
76 target.assign(&source);
77 }
78 }
79 }
80 }
81 let mut output = output.into_tensor();
82 unsafe { output.set_datum_type(input_tensor.datum_type()) }
83 Ok(output.into_tvalue())
84 }
85}
86
87impl Op for Pad {
88 fn name(&self) -> StaticName {
89 "Pad".into()
90 }
91
92 fn info(&self) -> TractResult<Vec<String>> {
93 Ok(vec![format!("Mode: {:?}, pads: {:?})", self.mode, self.pads,)])
94 }
95
96 op_as_typed_op!();
97}
98
99impl EvalOp for Pad {
100 op_out_of_plan!();
101
102 fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
103 let input = args_1!(inputs);
104 Ok(tvec!(dispatch_numbers!(Self::eval_t(input.datum_type())(self, input))?))
105 }
106}
107
108impl TypedOp for Pad {
109 as_op!();
110
111 fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
112 let mut fact = inputs[0].without_value();
113 if self.pads.len() != fact.rank() {
114 bail!("Inconsistent pad: input of rank {}, pads are: {:?}", fact.rank(), self.pads);
115 }
116 for (ix, (b, e)) in self.pads.iter().enumerate() {
117 fact.shape.set(ix, fact.shape[ix].clone() + *b + *e);
118 }
119 Ok(tvec!(fact))
120 }
121
122 fn input_roi(
123 &self,
124 model: &TypedModel,
125 node: &TypedNode,
126 ) -> TractResult<Option<TVec<Option<TDim>>>> {
127 let output_fact = model.outlet_fact(OutletId::new(node.id, 0))?;
128 rule_if_some!(roi = &output_fact.region_of_interest);
129 let mut input_roi = roi.clone();
131 for (axis, &(before, _)) in self.pads.iter().enumerate() {
132 if before == 0 {
133 continue;
134 }
135 if let Some(sym) = input_roi
136 .symbols()
137 .into_iter()
138 .find(|s| crate::ops::logic::sym_to_coord_axis(s) == Some(axis))
139 {
140 let shifted = TDim::Sym(sym.clone()) - TDim::Val(before as i64);
141 input_roi = input_roi.substitute(&sym, &shifted).unwrap_or(input_roi);
142 }
143 }
144 Ok(Some(tvec![Some(input_roi)]))
145 }
146
147 fn axes_mapping(
148 &self,
149 inputs: &[&TypedFact],
150 outputs: &[&TypedFact],
151 ) -> TractResult<AxesMapping> {
152 let mut result = AxesMapping::disconnected(inputs, outputs)?;
153 for (ix, pads) in self.pads.iter().enumerate() {
154 if pads == &(0, 0) {
155 result = result.linking((InOut::In(0), ix), (InOut::Out(0), ix))?;
156 }
157 }
158 Ok(result)
159 }
160
161 fn change_axes(
162 &self,
163 model: &TypedModel,
164 node: &TypedNode,
165 io: InOut,
166 change: &AxisOp,
167 ) -> TractResult<Option<AxisChangeConsequence>> {
168 let mut new_op = self.clone();
169 if let (InOut::In(0), AxisOp::Rm(ix)) = (io, change)
170 && new_op.pads.remove(*ix) == (0, 0)
171 {
172 return Ok(Some(AxisChangeConsequence::new(
173 model,
174 node,
175 Some(Box::new(new_op)),
176 change,
177 )));
178 }
179 if let (InOut::In(0), AxisOp::Add(ix)) = (io, change) {
180 new_op.pads.insert(*ix, (0, 0));
181 return Ok(Some(AxisChangeConsequence::new(
182 model,
183 node,
184 Some(Box::new(new_op)),
185 change,
186 )));
187 }
188 Ok(None)
189 }
190
191 fn declutter(
192 &self,
193 model: &TypedModel,
194 node: &TypedNode,
195 ) -> TractResult<Option<TypedModelPatch>> {
196 if self.pads.iter().all(|p| p.0 == 0 && p.1 == 0) {
197 TypedModelPatch::shunt_one_op(model, node)
198 } else {
199 Ok(None)
200 }
201 }
202}