1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
use tract_core::ndarray::*;
use tract_core::ops::array::PadMode;
use tract_nnef::internal::*;
#[derive(Debug, Clone, Default, Hash)]
struct PulsePadOpState {
current_pos: usize,
last_valid_frame: Option<Tensor>,
}
impl OpState for PulsePadOpState {
fn eval(
&mut self,
session: &mut SessionState,
op: &dyn Op,
mut inputs: TVec<Arc<Tensor>>,
) -> TractResult<TVec<Arc<Tensor>>> {
let input = args_1!(inputs).into_tensor();
let op = op.downcast_ref::<PulsePad>().ok_or_else(|| format_err!("Wrong Op type"))?;
let tensor = self.pad(session, op, input)?;
Ok(tvec!(tensor.into_arc_tensor()))
}
}
impl PulsePadOpState {
unsafe fn save_frame<T: Datum + Copy>(&mut self, op: &PulsePad, input: &Tensor, frame: usize) {
let data = input.to_array_view_unchecked::<T>();
self.last_valid_frame =
Some(data.index_axis(Axis(op.axis), frame).to_owned().into_tensor());
}
unsafe fn fill_slice_constant<T: Datum + Copy>(
data: &mut Tensor,
constant: &Tensor,
axis: usize,
range: std::ops::Range<usize>,
) {
let c = constant.to_scalar_unchecked::<T>();
data.to_array_view_mut_unchecked::<T>().slice_axis_mut(Axis(axis), range.into()).fill(*c);
}
unsafe fn fill_slice_with_frame<T: Datum + Copy>(
data: &mut Tensor,
axis: usize,
valid: &Tensor,
range: std::ops::Range<usize>,
) {
let mut data = data.to_array_view_mut_unchecked::<T>();
let valid = valid.to_array_view_unchecked::<T>();
for i in range {
data.slice_axis_mut(Axis(axis), (i..i + 1).into()).assign(&valid);
}
}
fn pad(
&mut self,
session: &mut SessionState,
op: &PulsePad,
mut input: Tensor,
) -> TractResult<Tensor> {
let pulse = input.shape()[op.axis];
let pulse_begin = self.current_pos;
let pulse_end = self.current_pos + pulse;
self.current_pos += pulse - op.overlap;
let end_input =
op.end_input.eval(&session.resolved_symbols).to_usize().unwrap_or(std::usize::MAX);
let after = op.after.eval(&session.resolved_symbols).to_usize().unwrap_or(std::usize::MAX);
if let PadMode::Edge = op.mode {
if after != 0 && pulse_begin < end_input {
let latest_valid_frame = (end_input - pulse_begin).min(pulse) - 1;
unsafe {
dispatch_copy_by_size!(Self::save_frame(input.datum_type())(
self,
op,
&input,
latest_valid_frame
))
}
}
}
if pulse_begin >= op.begin_input && pulse_end <= end_input {
return Ok(input);
}
if pulse_end <= op.begin_input - op.before || pulse_begin >= end_input.saturating_add(after)
{
return Ok(input);
}
if pulse_begin < op.begin_input {
let fill_up_to = (op.begin_input - pulse_begin).min(pulse);
match &op.mode {
PadMode::Constant(c) => unsafe {
dispatch_copy_by_size!(Self::fill_slice_constant(input.datum_type())(
&mut input,
c,
op.axis,
0..fill_up_to
))
},
PadMode::Edge => {
let frame = input.slice(op.axis, fill_up_to, fill_up_to + 1)?;
unsafe {
dispatch_copy_by_size!(Self::fill_slice_with_frame(input.datum_type())(
&mut input,
op.axis,
&frame,
0..fill_up_to
))
}
}
_ => unimplemented!(),
}
}
if pulse_end > end_input && after > 0 {
let fill_from = pulse - (pulse_end - end_input).min(pulse);
match &op.mode {
PadMode::Constant(c) => unsafe {
dispatch_copy_by_size!(Self::fill_slice_constant(input.datum_type())(
&mut input,
c,
op.axis,
fill_from..pulse
))
},
PadMode::Edge => {
let last_frame = self.last_valid_frame.as_ref().unwrap();
unsafe {
dispatch_copy_by_size!(Self::fill_slice_with_frame(input.datum_type())(
&mut input,
op.axis,
last_frame,
fill_from..pulse
))
}
}
_ => unimplemented!(),
}
}
Ok(input)
}
}
#[derive(Debug, Clone, Default, Hash)]
pub struct PulsePad {
pub axis: usize,
pub before: usize,
pub after: TDim,
pub begin_input: usize,
pub end_input: TDim,
pub mode: PadMode,
pub overlap: usize,
}
impl_dyn_hash!(PulsePad);
impl Op for PulsePad {
fn name(&self) -> Cow<str> {
"PulsePad".into()
}
fn info(&self) -> TractResult<Vec<String>> {
Ok(vec![format!(
"Mode: {:?}, axis: {} before: {} after: {}",
self.mode, self.axis, self.before, self.after,
)])
}
op_pulse!();
op_as_typed_op!();
}
impl EvalOp for PulsePad {
fn is_stateless(&self) -> bool {
false
}
fn state(
&self,
_session: &mut SessionState,
_node_id: usize,
) -> TractResult<Option<Box<dyn OpState>>> {
Ok(Some(Box::new(PulsePadOpState::default())))
}
}
impl TypedOp for PulsePad {
fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
Ok(tvec!(inputs[0].clone()))
}
as_op!();
}