1use super::*;
2
3impl ModuleOps<Self> for Dispatch {
4 fn batch_norm( x: FloatTensor<Self>, gamma: FloatTensor<Self>, beta: FloatTensor<Self>, mean: FloatTensor<Self>, variance: FloatTensor<Self>, epsilon: f64, ) -> FloatTensor<Self> {
5 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&gamma)).chain(core::iter::once(&beta)).chain(core::iter::once(&mean)).chain(core::iter::once(&variance)));
6 macro_rules! call {
7 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
8 let result = <$B>::batch_norm(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, beta, $context), unpack!($mode, $variant, Float, mean, $context), unpack!($mode, $variant, Float, variance, $context), epsilon);
9 pack!($mode, $variant, Float, result)
10 }};
11 }
12 dispatch!(selection, call)
13 }
14
15 fn batch_norm_train( x: FloatTensor<Self>, gamma: FloatTensor<Self>, beta: FloatTensor<Self>, epsilon: f64, ) -> BatchNormTrain<Self> {
16 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&gamma)).chain(core::iter::once(&beta)));
17 macro_rules! call {
18 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
19 let result = <$B>::batch_norm_train(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, beta, $context), epsilon);
20 BatchNormTrain { output: pack!($mode, $variant, Float, result.output), mean: pack!($mode, $variant, Float, result.mean), variance: pack!($mode, $variant, Float, result.variance) }
21 }};
22 }
23 dispatch!(selection, call)
24 }
25
26 fn batch_norm_train_backward( x: FloatTensor<Self>, gamma: FloatTensor<Self>, mean: FloatTensor<Self>, variance: FloatTensor<Self>, epsilon: f64, output_grad: FloatTensor<Self>, ) -> BatchNormTrainBackward<Self> {
27 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&gamma)).chain(core::iter::once(&mean)).chain(core::iter::once(&variance)).chain(core::iter::once(&output_grad)));
28 macro_rules! call {
29 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
30 let result = <$B>::batch_norm_train_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, mean, $context), unpack!($mode, $variant, Float, variance, $context), epsilon, unpack!($mode, $variant, Float, output_grad, $context));
31 BatchNormTrainBackward { x_grad: pack!($mode, $variant, Float, result.x_grad), gamma_grad: pack!($mode, $variant, Float, result.gamma_grad), beta_grad: pack!($mode, $variant, Float, result.beta_grad) }
32 }};
33 }
34 dispatch!(selection, call)
35 }
36
37 fn exponential_relu_native(tensor: FloatTensor<Self>, alpha: f64, continuous: bool) -> FloatTensor<Self> {
38 let selection = Selection::tensors(core::iter::once(&tensor));
39 macro_rules! call {
40 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
41 let result = <$B>::exponential_relu_native(unpack!($mode, $variant, Float, tensor, $context), alpha, continuous);
42 pack!($mode, $variant, Float, result)
43 }};
44 }
45 dispatch!(selection, call)
46 }
47
48 fn exponential_relu_native_backward(tensor: FloatTensor<Self>, grad: FloatTensor<Self>, alpha: f64, continuous: bool) -> FloatTensor<Self> {
49 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&grad)));
50 macro_rules! call {
51 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
52 let result = <$B>::exponential_relu_native_backward(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, grad, $context), alpha, continuous);
53 pack!($mode, $variant, Float, result)
54 }};
55 }
56 dispatch!(selection, call)
57 }
58
59 fn leaky_relu_native(tensor: FloatTensor<Self>, negative_slope: f64) -> FloatTensor<Self> {
60 let selection = Selection::tensors(core::iter::once(&tensor));
61 macro_rules! call {
62 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
63 let result = <$B>::leaky_relu_native(unpack!($mode, $variant, Float, tensor, $context), negative_slope);
64 pack!($mode, $variant, Float, result)
65 }};
66 }
67 dispatch!(selection, call)
68 }
69
70 fn leaky_relu_native_backward(tensor: FloatTensor<Self>, grad: FloatTensor<Self>, negative_slope: f64) -> FloatTensor<Self> {
71 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&grad)));
72 macro_rules! call {
73 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
74 let result = <$B>::leaky_relu_native_backward(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, grad, $context), negative_slope);
75 pack!($mode, $variant, Float, result)
76 }};
77 }
78 dispatch!(selection, call)
79 }
80
81 fn prelu_native(tensor: FloatTensor<Self>, alpha: FloatTensor<Self>) -> FloatTensor<Self> {
82 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&alpha)));
83 macro_rules! call {
84 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
85 let result = <$B>::prelu_native(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, alpha, $context));
86 pack!($mode, $variant, Float, result)
87 }};
88 }
89 dispatch!(selection, call)
90 }
91
92 fn prelu_native_backward_select(tensor: FloatTensor<Self>, alpha: FloatTensor<Self>, grad: FloatTensor<Self>, mask: [bool; 2]) -> [Option<FloatTensor<Self>>; 2] {
93 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&alpha)).chain(core::iter::once(&grad)));
94 macro_rules! call {
95 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
96 let result = <$B>::prelu_native_backward_select(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, alpha, $context), unpack!($mode, $variant, Float, grad, $context), mask);
97 result.map(|value| value.map(|value| pack!($mode, $variant, Float, value)))
98 }};
99 }
100 dispatch!(selection, call)
101 }
102
103 fn group_norm(tensor: FloatTensor<Self>, gamma: Option<FloatTensor<Self>>, beta: Option<FloatTensor<Self>>, groups: usize, epsilon: f64) -> FloatTensor<Self> {
104 let selection = Selection::tensors(core::iter::once(&tensor).chain(gamma.iter()).chain(beta.iter()));
105 macro_rules! call {
106 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
107 let result = <$B>::group_norm(unpack!($mode, $variant, Float, tensor, $context), gamma.map(|value| unpack!($mode, $variant, Float, value, $context)), beta.map(|value| unpack!($mode, $variant, Float, value, $context)), groups, epsilon);
108 pack!($mode, $variant, Float, result)
109 }};
110 }
111 dispatch!(selection, call)
112 }
113
114 fn group_norm_with_stats(tensor: FloatTensor<Self>, gamma: Option<FloatTensor<Self>>, beta: Option<FloatTensor<Self>>, groups: usize, epsilon: f64) -> LayerNormOutput<Self> {
115 let selection = Selection::tensors(core::iter::once(&tensor).chain(gamma.iter()).chain(beta.iter()));
116 macro_rules! call {
117 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
118 let result = <$B>::group_norm_with_stats(unpack!($mode, $variant, Float, tensor, $context), gamma.map(|value| unpack!($mode, $variant, Float, value, $context)), beta.map(|value| unpack!($mode, $variant, Float, value, $context)), groups, epsilon);
119 LayerNormOutput { output: pack!($mode, $variant, Float, result.output), mean: pack!($mode, $variant, Float, result.mean), rstd: pack!($mode, $variant, Float, result.rstd) }
120 }};
121 }
122 dispatch!(selection, call)
123 }
124
125 fn group_norm_backward_select(tensor: FloatTensor<Self>, gamma: Option<FloatTensor<Self>>, grad: FloatTensor<Self>, mean: FloatTensor<Self>, rstd: FloatTensor<Self>, groups: usize, mask: [bool; 3]) -> [Option<FloatTensor<Self>>; 3] {
126 let selection = Selection::tensors(core::iter::once(&tensor).chain(gamma.iter()).chain(core::iter::once(&grad)).chain(core::iter::once(&mean)).chain(core::iter::once(&rstd)));
127 macro_rules! call {
128 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
129 let result = <$B>::group_norm_backward_select(unpack!($mode, $variant, Float, tensor, $context), gamma.map(|value| unpack!($mode, $variant, Float, value, $context)), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Float, mean, $context), unpack!($mode, $variant, Float, rstd, $context), groups, mask);
130 result.map(|value| value.map(|value| pack!($mode, $variant, Float, value)))
131 }};
132 }
133 dispatch!(selection, call)
134 }
135
136 fn gelu_native(tensor: FloatTensor<Self>, approximate: bool) -> FloatTensor<Self> {
137 let selection = Selection::tensors(core::iter::once(&tensor));
138 macro_rules! call {
139 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
140 let result = <$B>::gelu_native(unpack!($mode, $variant, Float, tensor, $context), approximate);
141 pack!($mode, $variant, Float, result)
142 }};
143 }
144 dispatch!(selection, call)
145 }
146
147 fn gelu_native_backward(input: FloatTensor<Self>, grad: FloatTensor<Self>, approximate: bool) -> FloatTensor<Self> {
148 let selection = Selection::tensors(core::iter::once(&input).chain(core::iter::once(&grad)));
149 macro_rules! call {
150 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
151 let result = <$B>::gelu_native_backward(unpack!($mode, $variant, Float, input, $context), unpack!($mode, $variant, Float, grad, $context), approximate);
152 pack!($mode, $variant, Float, result)
153 }};
154 }
155 dispatch!(selection, call)
156 }
157
158 fn silu_native(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
159 let selection = Selection::tensors(core::iter::once(&tensor));
160 macro_rules! call {
161 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
162 let result = <$B>::silu_native(unpack!($mode, $variant, Float, tensor, $context));
163 pack!($mode, $variant, Float, result)
164 }};
165 }
166 dispatch!(selection, call)
167 }
168
169 fn silu_native_backward(input: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
170 let selection = Selection::tensors(core::iter::once(&input).chain(core::iter::once(&grad)));
171 macro_rules! call {
172 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
173 let result = <$B>::silu_native_backward(unpack!($mode, $variant, Float, input, $context), unpack!($mode, $variant, Float, grad, $context));
174 pack!($mode, $variant, Float, result)
175 }};
176 }
177 dispatch!(selection, call)
178 }
179
180 fn softmax_native(tensor: FloatTensor<Self>, dim: usize, logarithmic: bool) -> FloatTensor<Self> {
181 let selection = Selection::tensors(core::iter::once(&tensor));
182 macro_rules! call {
183 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
184 let result = <$B>::softmax_native(unpack!($mode, $variant, Float, tensor, $context), dim, logarithmic);
185 pack!($mode, $variant, Float, result)
186 }};
187 }
188 dispatch!(selection, call)
189 }
190
191 fn softmax_with_stats(tensor: FloatTensor<Self>, dim: usize, logarithmic: bool) -> SoftmaxOutput<Self> {
192 let selection = Selection::tensors(core::iter::once(&tensor));
193 macro_rules! call {
194 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
195 let result = <$B>::softmax_with_stats(unpack!($mode, $variant, Float, tensor, $context), dim, logarithmic);
196 SoftmaxOutput { output: pack!($mode, $variant, Float, result.output), working: pack!($mode, $variant, Float, result.working) }
197 }};
198 }
199 dispatch!(selection, call)
200 }
201
202 fn softmax_native_backward(working: FloatTensor<Self>, grad: FloatTensor<Self>, dim: usize, logarithmic: bool) -> FloatTensor<Self> {
203 let selection = Selection::tensors(core::iter::once(&working).chain(core::iter::once(&grad)));
204 macro_rules! call {
205 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
206 let result = <$B>::softmax_native_backward(unpack!($mode, $variant, Float, working, $context), unpack!($mode, $variant, Float, grad, $context), dim, logarithmic);
207 pack!($mode, $variant, Float, result)
208 }};
209 }
210 dispatch!(selection, call)
211 }
212
213 fn embedding(weights: FloatTensor<Self>, indices: IntTensor<Self>) -> FloatTensor<Self> {
214 let selection = Selection::tensors(core::iter::once(&weights).chain(core::iter::once(&indices)));
215 macro_rules! call {
216 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
217 let result = <$B>::embedding(unpack!($mode, $variant, Float, weights, $context), unpack!($mode, $variant, Int, indices, $context));
218 pack!($mode, $variant, Float, result)
219 }};
220 }
221 dispatch!(selection, call)
222 }
223
224 fn embedding_backward( weights: FloatTensor<Self>, output_grad: FloatTensor<Self>, indices: IntTensor<Self>, ) -> FloatTensor<Self> {
225 let selection = Selection::tensors(core::iter::once(&weights).chain(core::iter::once(&output_grad)).chain(core::iter::once(&indices)));
226 macro_rules! call {
227 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
228 let result = <$B>::embedding_backward(unpack!($mode, $variant, Float, weights, $context), unpack!($mode, $variant, Float, output_grad, $context), unpack!($mode, $variant, Int, indices, $context));
229 pack!($mode, $variant, Float, result)
230 }};
231 }
232 dispatch!(selection, call)
233 }
234
235 fn linear( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, ) -> FloatTensor<Self> {
236 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
237 macro_rules! call {
238 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
239 let result = <$B>::linear(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)));
240 pack!($mode, $variant, Float, result)
241 }};
242 }
243 dispatch!(selection, call)
244 }
245
246 fn linear_x_backward(weight: FloatTensor<Self>, output_grad: FloatTensor<Self>) -> FloatTensor<Self> {
247 let selection = Selection::tensors(core::iter::once(&weight).chain(core::iter::once(&output_grad)));
248 macro_rules! call {
249 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
250 let result = <$B>::linear_x_backward(unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context));
251 pack!($mode, $variant, Float, result)
252 }};
253 }
254 dispatch!(selection, call)
255 }
256
257 fn linear_weight_backward(x: FloatTensor<Self>, output_grad: FloatTensor<Self>) -> FloatTensor<Self> {
258 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&output_grad)));
259 macro_rules! call {
260 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
261 let result = <$B>::linear_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, output_grad, $context));
262 pack!($mode, $variant, Float, result)
263 }};
264 }
265 dispatch!(selection, call)
266 }
267
268 fn linear_bias_backward(output_grad: FloatTensor<Self>) -> FloatTensor<Self> {
269 let selection = Selection::tensors(core::iter::once(&output_grad));
270 macro_rules! call {
271 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
272 let result = <$B>::linear_bias_backward(unpack!($mode, $variant, Float, output_grad, $context));
273 pack!($mode, $variant, Float, result)
274 }};
275 }
276 dispatch!(selection, call)
277 }
278
279 fn conv1d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvOptions<1>, ) -> FloatTensor<Self> {
280 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
281 macro_rules! call {
282 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
283 let result = <$B>::conv1d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
284 pack!($mode, $variant, Float, result)
285 }};
286 }
287 dispatch!(selection, call)
288 }
289
290 fn conv1d_x_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<1>, ) -> FloatTensor<Self> {
291 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
292 macro_rules! call {
293 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
294 let result = <$B>::conv1d_x_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
295 pack!($mode, $variant, Float, result)
296 }};
297 }
298 dispatch!(selection, call)
299 }
300
301 fn conv1d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<1>, ) -> FloatTensor<Self> {
302 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
303 macro_rules! call {
304 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
305 let result = <$B>::conv1d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
306 pack!($mode, $variant, Float, result)
307 }};
308 }
309 dispatch!(selection, call)
310 }
311
312 fn conv1d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
313 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
314 macro_rules! call {
315 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
316 let result = <$B>::conv1d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
317 pack!($mode, $variant, Float, result)
318 }};
319 }
320 dispatch!(selection, call)
321 }
322
323 fn conv2d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvOptions<2>, ) -> FloatTensor<Self> {
324 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
325 macro_rules! call {
326 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
327 let result = <$B>::conv2d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
328 pack!($mode, $variant, Float, result)
329 }};
330 }
331 dispatch!(selection, call)
332 }
333
334 fn conv2d_x_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<2>, ) -> FloatTensor<Self> {
335 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
336 macro_rules! call {
337 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
338 let result = <$B>::conv2d_x_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
339 pack!($mode, $variant, Float, result)
340 }};
341 }
342 dispatch!(selection, call)
343 }
344
345 fn conv2d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<2>, ) -> FloatTensor<Self> {
346 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
347 macro_rules! call {
348 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
349 let result = <$B>::conv2d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
350 pack!($mode, $variant, Float, result)
351 }};
352 }
353 dispatch!(selection, call)
354 }
355
356 fn conv2d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
357 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
358 macro_rules! call {
359 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
360 let result = <$B>::conv2d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
361 pack!($mode, $variant, Float, result)
362 }};
363 }
364 dispatch!(selection, call)
365 }
366
367 fn deform_conv2d( x: FloatTensor<Self>, offset: FloatTensor<Self>, weight: FloatTensor<Self>, mask: Option<FloatTensor<Self>>, bias: Option<FloatTensor<Self>>, options: DeformConvOptions<2>, ) -> FloatTensor<Self> {
368 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&offset)).chain(core::iter::once(&weight)).chain(mask.iter()).chain(bias.iter()));
369 macro_rules! call {
370 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
371 let result = <$B>::deform_conv2d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, offset, $context), unpack!($mode, $variant, Float, weight, $context), mask.map(|value| unpack!($mode, $variant, Float, value, $context)), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
372 pack!($mode, $variant, Float, result)
373 }};
374 }
375 dispatch!(selection, call)
376 }
377
378 fn deform_conv2d_backward( x: FloatTensor<Self>, offset: FloatTensor<Self>, weight: FloatTensor<Self>, mask: Option<FloatTensor<Self>>, bias: Option<FloatTensor<Self>>, output_grad: FloatTensor<Self>, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<Self> {
379 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&offset)).chain(core::iter::once(&weight)).chain(mask.iter()).chain(bias.iter()).chain(core::iter::once(&output_grad)));
380 macro_rules! call {
381 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
382 let result = <$B>::deform_conv2d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, offset, $context), unpack!($mode, $variant, Float, weight, $context), mask.map(|value| unpack!($mode, $variant, Float, value, $context)), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), unpack!($mode, $variant, Float, output_grad, $context), options);
383 DeformConv2dBackward { x_grad: pack!($mode, $variant, Float, result.x_grad), offset_grad: pack!($mode, $variant, Float, result.offset_grad), weight_grad: pack!($mode, $variant, Float, result.weight_grad), mask_grad: result.mask_grad.map(|value| pack!($mode, $variant, Float, value)), bias_grad: result.bias_grad.map(|value| pack!($mode, $variant, Float, value)) }
384 }};
385 }
386 dispatch!(selection, call)
387 }
388
389 fn conv3d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvOptions<3>, ) -> FloatTensor<Self> {
390 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
391 macro_rules! call {
392 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
393 let result = <$B>::conv3d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
394 pack!($mode, $variant, Float, result)
395 }};
396 }
397 dispatch!(selection, call)
398 }
399
400 fn conv3d_x_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<3>, ) -> FloatTensor<Self> {
401 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
402 macro_rules! call {
403 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
404 let result = <$B>::conv3d_x_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
405 pack!($mode, $variant, Float, result)
406 }};
407 }
408 dispatch!(selection, call)
409 }
410
411 fn conv3d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvOptions<3>, ) -> FloatTensor<Self> {
412 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
413 macro_rules! call {
414 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
415 let result = <$B>::conv3d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
416 pack!($mode, $variant, Float, result)
417 }};
418 }
419 dispatch!(selection, call)
420 }
421
422 fn conv3d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
423 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
424 macro_rules! call {
425 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
426 let result = <$B>::conv3d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
427 pack!($mode, $variant, Float, result)
428 }};
429 }
430 dispatch!(selection, call)
431 }
432
433 fn conv_transpose1d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvTransposeOptions<1>, ) -> FloatTensor<Self> {
434 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
435 macro_rules! call {
436 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
437 let result = <$B>::conv_transpose1d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
438 pack!($mode, $variant, Float, result)
439 }};
440 }
441 dispatch!(selection, call)
442 }
443
444 fn conv_transpose1d_x_backward( weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<1>, ) -> FloatTensor<Self> {
445 let selection = Selection::tensors(core::iter::once(&weight).chain(core::iter::once(&output_grad)));
446 macro_rules! call {
447 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
448 let result = <$B>::conv_transpose1d_x_backward(unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
449 pack!($mode, $variant, Float, result)
450 }};
451 }
452 dispatch!(selection, call)
453 }
454
455 fn conv_transpose1d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<1>, ) -> FloatTensor<Self> {
456 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
457 macro_rules! call {
458 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
459 let result = <$B>::conv_transpose1d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
460 pack!($mode, $variant, Float, result)
461 }};
462 }
463 dispatch!(selection, call)
464 }
465
466 fn conv_transpose1d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
467 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
468 macro_rules! call {
469 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
470 let result = <$B>::conv_transpose1d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
471 pack!($mode, $variant, Float, result)
472 }};
473 }
474 dispatch!(selection, call)
475 }
476
477 fn conv_transpose2d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvTransposeOptions<2>, ) -> FloatTensor<Self> {
478 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
479 macro_rules! call {
480 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
481 let result = <$B>::conv_transpose2d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
482 pack!($mode, $variant, Float, result)
483 }};
484 }
485 dispatch!(selection, call)
486 }
487
488 fn conv_transpose2d_x_backward( weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<2>, ) -> FloatTensor<Self> {
489 let selection = Selection::tensors(core::iter::once(&weight).chain(core::iter::once(&output_grad)));
490 macro_rules! call {
491 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
492 let result = <$B>::conv_transpose2d_x_backward(unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
493 pack!($mode, $variant, Float, result)
494 }};
495 }
496 dispatch!(selection, call)
497 }
498
499 fn conv_transpose2d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<2>, ) -> FloatTensor<Self> {
500 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
501 macro_rules! call {
502 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
503 let result = <$B>::conv_transpose2d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
504 pack!($mode, $variant, Float, result)
505 }};
506 }
507 dispatch!(selection, call)
508 }
509
510 fn conv_transpose2d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
511 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
512 macro_rules! call {
513 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
514 let result = <$B>::conv_transpose2d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
515 pack!($mode, $variant, Float, result)
516 }};
517 }
518 dispatch!(selection, call)
519 }
520
521 fn conv_transpose3d( x: FloatTensor<Self>, weight: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, options: ConvTransposeOptions<3>, ) -> FloatTensor<Self> {
522 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(bias.iter()));
523 macro_rules! call {
524 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
525 let result = <$B>::conv_transpose3d(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
526 pack!($mode, $variant, Float, result)
527 }};
528 }
529 dispatch!(selection, call)
530 }
531
532 fn conv_transpose3d_x_backward( weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<3>, ) -> FloatTensor<Self> {
533 let selection = Selection::tensors(core::iter::once(&weight).chain(core::iter::once(&output_grad)));
534 macro_rules! call {
535 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
536 let result = <$B>::conv_transpose3d_x_backward(unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
537 pack!($mode, $variant, Float, result)
538 }};
539 }
540 dispatch!(selection, call)
541 }
542
543 fn conv_transpose3d_weight_backward( x: FloatTensor<Self>, weight: FloatTensor<Self>, output_grad: FloatTensor<Self>, options: ConvTransposeOptions<3>, ) -> FloatTensor<Self> {
544 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&weight)).chain(core::iter::once(&output_grad)));
545 macro_rules! call {
546 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
547 let result = <$B>::conv_transpose3d_weight_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, weight, $context), unpack!($mode, $variant, Float, output_grad, $context), options);
548 pack!($mode, $variant, Float, result)
549 }};
550 }
551 dispatch!(selection, call)
552 }
553
554 fn conv_transpose3d_bias_backward( x: FloatTensor<Self>, bias: FloatTensor<Self>, output_grad: FloatTensor<Self>, ) -> FloatTensor<Self> {
555 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&bias)).chain(core::iter::once(&output_grad)));
556 macro_rules! call {
557 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
558 let result = <$B>::conv_transpose3d_bias_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, bias, $context), unpack!($mode, $variant, Float, output_grad, $context));
559 pack!($mode, $variant, Float, result)
560 }};
561 }
562 dispatch!(selection, call)
563 }
564
565 fn unfold4d( x: FloatTensor<Self>, kernel_size: [usize; 2], options: UnfoldOptions, ) -> FloatTensor<Self> {
566 let selection = Selection::tensors(core::iter::once(&x));
567 macro_rules! call {
568 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
569 let result = <$B>::unfold4d(unpack!($mode, $variant, Float, x, $context), kernel_size, options);
570 pack!($mode, $variant, Float, result)
571 }};
572 }
573 dispatch!(selection, call)
574 }
575
576 fn fold4d( x: FloatTensor<Self>, output_size: [usize; 2], kernel_size: [usize; 2], options: UnfoldOptions, ) -> FloatTensor<Self> {
577 let selection = Selection::tensors(core::iter::once(&x));
578 macro_rules! call {
579 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
580 let result = <$B>::fold4d(unpack!($mode, $variant, Float, x, $context), output_size, kernel_size, options);
581 pack!($mode, $variant, Float, result)
582 }};
583 }
584 dispatch!(selection, call)
585 }
586
587 fn avg_pool1d( x: FloatTensor<Self>, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<Self> {
588 let selection = Selection::tensors(core::iter::once(&x));
589 macro_rules! call {
590 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
591 let result = <$B>::avg_pool1d(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
592 pack!($mode, $variant, Float, result)
593 }};
594 }
595 dispatch!(selection, call)
596 }
597
598 fn avg_pool1d_backward( x: FloatTensor<Self>, grad: FloatTensor<Self>, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<Self> {
599 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
600 macro_rules! call {
601 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
602 let result = <$B>::avg_pool1d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
603 pack!($mode, $variant, Float, result)
604 }};
605 }
606 dispatch!(selection, call)
607 }
608
609 fn avg_pool2d( x: FloatTensor<Self>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<Self> {
610 let selection = Selection::tensors(core::iter::once(&x));
611 macro_rules! call {
612 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
613 let result = <$B>::avg_pool2d(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
614 pack!($mode, $variant, Float, result)
615 }};
616 }
617 dispatch!(selection, call)
618 }
619
620 fn avg_pool2d_backward( x: FloatTensor<Self>, grad: FloatTensor<Self>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<Self> {
621 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
622 macro_rules! call {
623 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
624 let result = <$B>::avg_pool2d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
625 pack!($mode, $variant, Float, result)
626 }};
627 }
628 dispatch!(selection, call)
629 }
630
631 fn avg_pool3d(x: FloatTensor<Self>, kernel_size: [usize; 3], stride: [usize; 3], padding: [usize; 3], count_include_pad: bool, ceil_mode: bool) -> FloatTensor<Self> {
632 let selection = Selection::tensors(core::iter::once(&x));
633 macro_rules! call {
634 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
635 let result = <$B>::avg_pool3d(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
636 pack!($mode, $variant, Float, result)
637 }};
638 }
639 dispatch!(selection, call)
640 }
641
642 fn avg_pool3d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>, kernel_size: [usize; 3], stride: [usize; 3], padding: [usize; 3], count_include_pad: bool, ceil_mode: bool) -> FloatTensor<Self> {
643 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
644 macro_rules! call {
645 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
646 let result = <$B>::avg_pool3d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), kernel_size, stride, padding, count_include_pad, ceil_mode);
647 pack!($mode, $variant, Float, result)
648 }};
649 }
650 dispatch!(selection, call)
651 }
652
653 fn adaptive_avg_pool2d(x: FloatTensor<Self>, output_size: [usize; 2]) -> FloatTensor<Self> {
654 let selection = Selection::tensors(core::iter::once(&x));
655 macro_rules! call {
656 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
657 let result = <$B>::adaptive_avg_pool2d(unpack!($mode, $variant, Float, x, $context), output_size);
658 pack!($mode, $variant, Float, result)
659 }};
660 }
661 dispatch!(selection, call)
662 }
663
664 fn adaptive_avg_pool2d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
665 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
666 macro_rules! call {
667 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
668 let result = <$B>::adaptive_avg_pool2d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context));
669 pack!($mode, $variant, Float, result)
670 }};
671 }
672 dispatch!(selection, call)
673 }
674
675 fn adaptive_avg_pool3d(x: FloatTensor<Self>, output_size: [usize; 3]) -> FloatTensor<Self> {
676 let selection = Selection::tensors(core::iter::once(&x));
677 macro_rules! call {
678 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
679 let result = <$B>::adaptive_avg_pool3d(unpack!($mode, $variant, Float, x, $context), output_size);
680 pack!($mode, $variant, Float, result)
681 }};
682 }
683 dispatch!(selection, call)
684 }
685
686 fn adaptive_avg_pool3d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
687 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
688 macro_rules! call {
689 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
690 let result = <$B>::adaptive_avg_pool3d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context));
691 pack!($mode, $variant, Float, result)
692 }};
693 }
694 dispatch!(selection, call)
695 }
696
697 fn max_pool3d(x: FloatTensor<Self>, kernel: [usize; 3], stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], ceil: bool) -> FloatTensor<Self> {
698 let selection = Selection::tensors(core::iter::once(&x));
699 macro_rules! call {
700 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
701 let result = <$B>::max_pool3d(unpack!($mode, $variant, Float, x, $context), kernel, stride, padding, dilation, ceil);
702 pack!($mode, $variant, Float, result)
703 }};
704 }
705 dispatch!(selection, call)
706 }
707
708 fn max_pool3d_with_indices(x: FloatTensor<Self>, kernel: [usize; 3], stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], ceil: bool) -> MaxPool3dWithIndices<Self> {
709 let selection = Selection::tensors(core::iter::once(&x));
710 macro_rules! call {
711 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
712 let result = <$B>::max_pool3d_with_indices(unpack!($mode, $variant, Float, x, $context), kernel, stride, padding, dilation, ceil);
713 MaxPool3dWithIndices { output: pack!($mode, $variant, Float, result.output), indices: pack!($mode, $variant, Int, result.indices) }
714 }};
715 }
716 dispatch!(selection, call)
717 }
718
719 fn max_pool3d_with_indices_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>, indices: IntTensor<Self>, _kernel: [usize; 3], _stride: [usize; 3], _padding: [usize; 3], _dilation: [usize; 3], _ceil: bool) -> MaxPool3dBackward<Self> {
720 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)).chain(core::iter::once(&indices)));
721 macro_rules! call {
722 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
723 let result = <$B>::max_pool3d_with_indices_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Int, indices, $context), _kernel, _stride, _padding, _dilation, _ceil);
724 MaxPool3dBackward { x_grad: pack!($mode, $variant, Float, result.x_grad) }
725 }};
726 }
727 dispatch!(selection, call)
728 }
729
730 fn adaptive_avg_pool1d(x: FloatTensor<Self>, output_size: usize) -> FloatTensor<Self> {
731 let selection = Selection::tensors(core::iter::once(&x));
732 macro_rules! call {
733 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
734 let result = <$B>::adaptive_avg_pool1d(unpack!($mode, $variant, Float, x, $context), output_size);
735 pack!($mode, $variant, Float, result)
736 }};
737 }
738 dispatch!(selection, call)
739 }
740
741 fn adaptive_avg_pool1d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
742 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
743 macro_rules! call {
744 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
745 let result = <$B>::adaptive_avg_pool1d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context));
746 pack!($mode, $variant, Float, result)
747 }};
748 }
749 dispatch!(selection, call)
750 }
751
752 fn max_pool1d( x: FloatTensor<Self>, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, ) -> FloatTensor<Self> {
753 let selection = Selection::tensors(core::iter::once(&x));
754 macro_rules! call {
755 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
756 let result = <$B>::max_pool1d(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode);
757 pack!($mode, $variant, Float, result)
758 }};
759 }
760 dispatch!(selection, call)
761 }
762
763 fn max_pool1d_with_indices( x: FloatTensor<Self>, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, ) -> MaxPool1dWithIndices<Self> {
764 let selection = Selection::tensors(core::iter::once(&x));
765 macro_rules! call {
766 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
767 let result = <$B>::max_pool1d_with_indices(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode);
768 MaxPool1dWithIndices { output: pack!($mode, $variant, Float, result.output), indices: pack!($mode, $variant, Int, result.indices) }
769 }};
770 }
771 dispatch!(selection, call)
772 }
773
774 fn max_pool1d_with_indices_backward( x: FloatTensor<Self>, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, output_grad: FloatTensor<Self>, indices: IntTensor<Self>, ) -> MaxPool1dBackward<Self> {
775 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&output_grad)).chain(core::iter::once(&indices)));
776 macro_rules! call {
777 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
778 let result = <$B>::max_pool1d_with_indices_backward(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode, unpack!($mode, $variant, Float, output_grad, $context), unpack!($mode, $variant, Int, indices, $context));
779 MaxPool1dBackward { x_grad: pack!($mode, $variant, Float, result.x_grad) }
780 }};
781 }
782 dispatch!(selection, call)
783 }
784
785 fn max_pool2d( x: FloatTensor<Self>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> FloatTensor<Self> {
786 let selection = Selection::tensors(core::iter::once(&x));
787 macro_rules! call {
788 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
789 let result = <$B>::max_pool2d(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode);
790 pack!($mode, $variant, Float, result)
791 }};
792 }
793 dispatch!(selection, call)
794 }
795
796 fn max_pool2d_with_indices( x: FloatTensor<Self>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> MaxPool2dWithIndices<Self> {
797 let selection = Selection::tensors(core::iter::once(&x));
798 macro_rules! call {
799 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
800 let result = <$B>::max_pool2d_with_indices(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode);
801 MaxPool2dWithIndices { output: pack!($mode, $variant, Float, result.output), indices: pack!($mode, $variant, Int, result.indices) }
802 }};
803 }
804 dispatch!(selection, call)
805 }
806
807 fn max_pool2d_with_indices_backward( x: FloatTensor<Self>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, output_grad: FloatTensor<Self>, indices: IntTensor<Self>, ) -> MaxPool2dBackward<Self> {
808 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&output_grad)).chain(core::iter::once(&indices)));
809 macro_rules! call {
810 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
811 let result = <$B>::max_pool2d_with_indices_backward(unpack!($mode, $variant, Float, x, $context), kernel_size, stride, padding, dilation, ceil_mode, unpack!($mode, $variant, Float, output_grad, $context), unpack!($mode, $variant, Int, indices, $context));
812 MaxPool2dBackward { x_grad: pack!($mode, $variant, Float, result.x_grad) }
813 }};
814 }
815 dispatch!(selection, call)
816 }
817
818 fn interpolate( x: FloatTensor<Self>, output_size: [usize; 2], options: InterpolateOptions, ) -> FloatTensor<Self> {
819 let selection = Selection::tensors(core::iter::once(&x));
820 macro_rules! call {
821 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
822 let result = <$B>::interpolate(unpack!($mode, $variant, Float, x, $context), output_size, options);
823 pack!($mode, $variant, Float, result)
824 }};
825 }
826 dispatch!(selection, call)
827 }
828
829 fn interpolate_backward( x: FloatTensor<Self>, grad: FloatTensor<Self>, output_size: [usize; 2], options: InterpolateOptions, ) -> FloatTensor<Self> {
830 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
831 macro_rules! call {
832 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
833 let result = <$B>::interpolate_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), output_size, options);
834 pack!($mode, $variant, Float, result)
835 }};
836 }
837 dispatch!(selection, call)
838 }
839
840 fn interpolate1d(x: FloatTensor<Self>, size: usize, options: InterpolateOptions) -> FloatTensor<Self> {
841 let selection = Selection::tensors(core::iter::once(&x));
842 macro_rules! call {
843 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
844 let result = <$B>::interpolate1d(unpack!($mode, $variant, Float, x, $context), size, options);
845 pack!($mode, $variant, Float, result)
846 }};
847 }
848 dispatch!(selection, call)
849 }
850
851 fn interpolate1d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>, size: usize, options: InterpolateOptions) -> FloatTensor<Self> {
852 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
853 macro_rules! call {
854 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
855 let result = <$B>::interpolate1d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), size, options);
856 pack!($mode, $variant, Float, result)
857 }};
858 }
859 dispatch!(selection, call)
860 }
861
862 fn interpolate3d(x: FloatTensor<Self>, size: [usize; 3], options: InterpolateOptions) -> FloatTensor<Self> {
863 let selection = Selection::tensors(core::iter::once(&x));
864 macro_rules! call {
865 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
866 let result = <$B>::interpolate3d(unpack!($mode, $variant, Float, x, $context), size, options);
867 pack!($mode, $variant, Float, result)
868 }};
869 }
870 dispatch!(selection, call)
871 }
872
873 fn interpolate3d_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>, size: [usize; 3], options: InterpolateOptions) -> FloatTensor<Self> {
874 let selection = Selection::tensors(core::iter::once(&x).chain(core::iter::once(&grad)));
875 macro_rules! call {
876 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
877 let result = <$B>::interpolate3d_backward(unpack!($mode, $variant, Float, x, $context), unpack!($mode, $variant, Float, grad, $context), size, options);
878 pack!($mode, $variant, Float, result)
879 }};
880 }
881 dispatch!(selection, call)
882 }
883
884 fn attention( query: FloatTensor<Self>, key: FloatTensor<Self>, value: FloatTensor<Self>, mask: Option<BoolTensor<Self>>, attn_bias: Option<FloatTensor<Self>>, options: AttentionModuleOptions, ) -> FloatTensor<Self> {
885 let selection = Selection::tensors(core::iter::once(&query).chain(core::iter::once(&key)).chain(core::iter::once(&value)).chain(mask.iter()).chain(attn_bias.iter()));
886 macro_rules! call {
887 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
888 let result = <$B>::attention(unpack!($mode, $variant, Float, query, $context), unpack!($mode, $variant, Float, key, $context), unpack!($mode, $variant, Float, value, $context), mask.map(|value| unpack!($mode, $variant, Bool, value, $context)), attn_bias.map(|value| unpack!($mode, $variant, Float, value, $context)), options);
889 pack!($mode, $variant, Float, result)
890 }};
891 }
892 dispatch!(selection, call)
893 }
894
895 fn layer_norm( tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, beta: Option<FloatTensor<Self>>, epsilon: f64, ) -> FloatTensor<Self> {
896 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(beta.iter()));
897 macro_rules! call {
898 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
899 let result = <$B>::layer_norm(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), beta.map(|value| unpack!($mode, $variant, Float, value, $context)), epsilon);
900 pack!($mode, $variant, Float, result)
901 }};
902 }
903 dispatch!(selection, call)
904 }
905
906 fn rms_norm(tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, epsilon: f64) -> FloatTensor<Self> {
907 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)));
908 macro_rules! call {
909 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
910 let result = <$B>::rms_norm(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), epsilon);
911 pack!($mode, $variant, Float, result)
912 }};
913 }
914 dispatch!(selection, call)
915 }
916
917 fn rms_norm_with_stats(tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, epsilon: f64) -> RmsNormOutput<Self> {
918 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)));
919 macro_rules! call {
920 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
921 let result = <$B>::rms_norm_with_stats(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), epsilon);
922 RmsNormOutput { output: pack!($mode, $variant, Float, result.output), rstd: pack!($mode, $variant, Float, result.rstd) }
923 }};
924 }
925 dispatch!(selection, call)
926 }
927
928 fn rms_norm_backward(tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, grad: FloatTensor<Self>, rstd: FloatTensor<Self>) -> RmsNormBackward<Self> {
929 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(core::iter::once(&grad)).chain(core::iter::once(&rstd)));
930 macro_rules! call {
931 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
932 let result = <$B>::rms_norm_backward(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Float, rstd, $context));
933 RmsNormBackward { input: pack!($mode, $variant, Float, result.input), weight: pack!($mode, $variant, Float, result.weight) }
934 }};
935 }
936 dispatch!(selection, call)
937 }
938
939 fn rms_norm_backward_select(tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, grad: FloatTensor<Self>, rstd: FloatTensor<Self>, mask: [bool; 2]) -> [Option<FloatTensor<Self>>; 2] {
940 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(core::iter::once(&grad)).chain(core::iter::once(&rstd)));
941 macro_rules! call {
942 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
943 let result = <$B>::rms_norm_backward_select(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Float, rstd, $context), mask);
944 result.map(|value| value.map(|value| pack!($mode, $variant, Float, value)))
945 }};
946 }
947 dispatch!(selection, call)
948 }
949
950 fn layer_norm_with_stats( tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, beta: Option<FloatTensor<Self>>, epsilon: f64, ) -> LayerNormOutput<Self> {
951 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(beta.iter()));
952 macro_rules! call {
953 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
954 let result = <$B>::layer_norm_with_stats(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), beta.map(|value| unpack!($mode, $variant, Float, value, $context)), epsilon);
955 LayerNormOutput { output: pack!($mode, $variant, Float, result.output), mean: pack!($mode, $variant, Float, result.mean), rstd: pack!($mode, $variant, Float, result.rstd) }
956 }};
957 }
958 dispatch!(selection, call)
959 }
960
961 fn layer_norm_backward( tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, grad: FloatTensor<Self>, mean: FloatTensor<Self>, rstd: FloatTensor<Self>, ) -> LayerNormBackward<Self> {
962 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(core::iter::once(&grad)).chain(core::iter::once(&mean)).chain(core::iter::once(&rstd)));
963 macro_rules! call {
964 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
965 let result = <$B>::layer_norm_backward(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Float, mean, $context), unpack!($mode, $variant, Float, rstd, $context));
966 LayerNormBackward { input: pack!($mode, $variant, Float, result.input), weight: pack!($mode, $variant, Float, result.weight), bias: pack!($mode, $variant, Float, result.bias) }
967 }};
968 }
969 dispatch!(selection, call)
970 }
971
972 fn layer_norm_backward_select(tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, grad: FloatTensor<Self>, mean: FloatTensor<Self>, rstd: FloatTensor<Self>, mask: [bool; 3]) -> [Option<FloatTensor<Self>>; 3] {
973 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(core::iter::once(&grad)).chain(core::iter::once(&mean)).chain(core::iter::once(&rstd)));
974 macro_rules! call {
975 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
976 let result = <$B>::layer_norm_backward_select(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), unpack!($mode, $variant, Float, grad, $context), unpack!($mode, $variant, Float, mean, $context), unpack!($mode, $variant, Float, rstd, $context), mask);
977 result.map(|value| value.map(|value| pack!($mode, $variant, Float, value)))
978 }};
979 }
980 dispatch!(selection, call)
981 }
982
983 fn layer_norm_default( tensor: FloatTensor<Self>, gamma: FloatTensor<Self>, beta: Option<FloatTensor<Self>>, epsilon: f64, ) -> FloatTensor<Self> {
984 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&gamma)).chain(beta.iter()));
985 macro_rules! call {
986 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
987 let result = <$B>::layer_norm_default(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, gamma, $context), beta.map(|value| unpack!($mode, $variant, Float, value, $context)), epsilon);
988 pack!($mode, $variant, Float, result)
989 }};
990 }
991 dispatch!(selection, call)
992 }
993
994 fn ctc_loss( log_probs: FloatTensor<Self>, targets: IntTensor<Self>, input_lengths: IntTensor<Self>, target_lengths: IntTensor<Self>, blank: usize, ) -> FloatTensor<Self> {
995 let selection = Selection::tensors(core::iter::once(&log_probs).chain(core::iter::once(&targets)).chain(core::iter::once(&input_lengths)).chain(core::iter::once(&target_lengths)));
996 macro_rules! call {
997 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
998 let result = <$B>::ctc_loss(unpack!($mode, $variant, Float, log_probs, $context), unpack!($mode, $variant, Int, targets, $context), unpack!($mode, $variant, Int, input_lengths, $context), unpack!($mode, $variant, Int, target_lengths, $context), blank);
999 pack!($mode, $variant, Float, result)
1000 }};
1001 }
1002 dispatch!(selection, call)
1003 }
1004
1005 fn ctc_loss_backward( _log_probs: FloatTensor<Self>, _targets: IntTensor<Self>, _input_lengths: IntTensor<Self>, _target_lengths: IntTensor<Self>, _grad_loss: FloatTensor<Self>, _blank: usize, ) -> FloatTensor<Self> {
1006 let selection = Selection::tensors(core::iter::once(&_log_probs).chain(core::iter::once(&_targets)).chain(core::iter::once(&_input_lengths)).chain(core::iter::once(&_target_lengths)).chain(core::iter::once(&_grad_loss)));
1007 macro_rules! call {
1008 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1009 let result = <$B>::ctc_loss_backward(unpack!($mode, $variant, Float, _log_probs, $context), unpack!($mode, $variant, Int, _targets, $context), unpack!($mode, $variant, Int, _input_lengths, $context), unpack!($mode, $variant, Int, _target_lengths, $context), unpack!($mode, $variant, Float, _grad_loss, $context), _blank);
1010 pack!($mode, $variant, Float, result)
1011 }};
1012 }
1013 dispatch!(selection, call)
1014 }
1015
1016 fn rfft( signal: FloatTensor<Self>, dim: usize, n: Option<usize>, ) -> (FloatTensor<Self>, FloatTensor<Self>) {
1017 let selection = Selection::tensors(core::iter::once(&signal));
1018 macro_rules! call {
1019 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1020 let result = <$B>::rfft(unpack!($mode, $variant, Float, signal, $context), dim, n);
1021 { let (value0, value1) = result; (pack!($mode, $variant, Float, value0), pack!($mode, $variant, Float, value1)) }
1022 }};
1023 }
1024 dispatch!(selection, call)
1025 }
1026
1027 fn irfft( spectrum_re: FloatTensor<Self>, spectrum_im: FloatTensor<Self>, dim: usize, n: Option<usize>, ) -> FloatTensor<Self> {
1028 let selection = Selection::tensors(core::iter::once(&spectrum_re).chain(core::iter::once(&spectrum_im)));
1029 macro_rules! call {
1030 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1031 let result = <$B>::irfft(unpack!($mode, $variant, Float, spectrum_re, $context), unpack!($mode, $variant, Float, spectrum_im, $context), dim, n);
1032 pack!($mode, $variant, Float, result)
1033 }};
1034 }
1035 dispatch!(selection, call)
1036 }
1037}