burn_backend/backend/ops/modules/base.rs
1use super::{conv, ctc, linear, pool};
2use crate::ops::unfold::{create_unfolding_weight, unfold4d_using_conv2d};
3use crate::tensor::{BoolTensor, FloatTensor, IntTensor};
4use crate::{Backend, Scalar, TensorMetadata};
5pub use burn_std::ops::{
6 AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConvOptions,
7 GridSampleOptions, GridSamplePaddingMode, InterpolateMode, InterpolateOptions, PadMode,
8 PaddedConvOptions, UnfoldOptions,
9};
10use burn_std::{IntDType, Shape};
11
12/// Gradient computed during the backward pass for each tensor used by [conv2d](ModuleOps::conv2d).
13#[derive(new)]
14pub struct Conv2dBackward<B: Backend> {
15 /// Gradient.
16 pub x_grad: FloatTensor<B>,
17
18 /// Weights gradient.
19 pub weights_grad: FloatTensor<B>,
20
21 /// Bias gradient.
22 pub bias_grad: Option<FloatTensor<B>>,
23}
24
25/// Gradient computed during the backward pass for each tensor used by [deform_conv2d](ModuleOps::deform_conv2d).
26#[derive(new)]
27pub struct DeformConv2dBackward<B: Backend> {
28 /// Gradient.
29 pub x_grad: FloatTensor<B>,
30
31 /// Offset gradient.
32 pub offset_grad: FloatTensor<B>,
33
34 /// Weights gradient.
35 pub weight_grad: FloatTensor<B>,
36
37 /// Mask gradient.
38 pub mask_grad: Option<FloatTensor<B>>,
39
40 /// Bias gradient.
41 pub bias_grad: Option<FloatTensor<B>>,
42}
43
44/// Gradient computed during the backward pass for each tensor used by [conv3d](ModuleOps::conv3d).
45#[derive(new)]
46pub struct Conv3dBackward<B: Backend> {
47 /// Gradient.
48 pub x_grad: FloatTensor<B>,
49
50 /// Weights gradient.
51 pub weights_grad: FloatTensor<B>,
52
53 /// Bias gradient.
54 pub bias_grad: Option<FloatTensor<B>>,
55}
56
57/// Gradient computed during the backward pass for each tensor used by [max_pool1d](ModuleOps::max_pool1d).
58#[derive(new)]
59pub struct MaxPool1dBackward<B: Backend> {
60 /// Gradient.
61 pub x_grad: FloatTensor<B>,
62}
63
64/// Results from [max_pool1d](ModuleOps::max_pool1d_with_indices).
65#[derive(new)]
66pub struct MaxPool1dWithIndices<B: Backend> {
67 /// The output tensor.
68 pub output: FloatTensor<B>,
69
70 /// The indices tensor.
71 pub indices: IntTensor<B>,
72}
73
74/// Gradient computed during the backward pass for each tensor used by [max_pool2d](ModuleOps::max_pool2d).
75#[derive(new)]
76pub struct MaxPool2dBackward<B: Backend> {
77 /// Gradient.
78 pub x_grad: FloatTensor<B>,
79}
80
81/// Results from [max_pool2d](ModuleOps::max_pool2d_with_indices).
82#[derive(new)]
83pub struct MaxPool2dWithIndices<B: Backend> {
84 /// The output tensor.
85 pub output: FloatTensor<B>,
86
87 /// The indices tensor.
88 pub indices: IntTensor<B>,
89}
90
91/// Gradient computed during the backward pass for each tensor used by [interpolate](ModuleOps::interpolate).
92#[derive(new)]
93pub struct InterpolateBackward<B: Backend> {
94 /// Gradient.
95 pub x_grad: FloatTensor<B>,
96}
97
98/// Module operations trait.
99pub trait ModuleOps<B: Backend> {
100 /// Applies batch normalization using explicitly supplied channel statistics.
101 ///
102 /// The input has shape `[batch, channels, ...]`; all other tensors have
103 /// shape `[channels]`.
104 ///
105 /// This operation doesn't calculate or update statistics. Callers may
106 /// supply running statistics for inference or batch statistics calculated
107 /// by a training path.
108 fn batch_norm(
109 x: FloatTensor<B>,
110 gamma: FloatTensor<B>,
111 beta: FloatTensor<B>,
112 mean: FloatTensor<B>,
113 variance: FloatTensor<B>,
114 epsilon: f64,
115 ) -> FloatTensor<B> {
116 let rank = x.shape().num_dims();
117 let channels = x.shape()[1];
118 let mut dimensions = alloc::vec![1; rank];
119 dimensions[1] = channels;
120 let shape = Shape::from(dimensions);
121 let gamma = B::float_reshape(gamma, shape.clone());
122 let beta = B::float_reshape(beta, shape.clone());
123 let mean = B::float_reshape(mean, shape.clone());
124 let variance = B::float_reshape(variance, shape);
125 let std = B::float_sqrt(B::float_add_scalar(variance, Scalar::Float(epsilon)));
126 let normalized = B::float_div(B::float_sub(x, mean), std);
127 B::float_add(B::float_mul(normalized, gamma), beta)
128 }
129
130 /// Embedding operation.
131 ///
132 /// # Arguments
133 ///
134 /// * `weights` - The embedding weights.
135 /// * `indices` - The indices tensor.
136 ///
137 /// # Returns
138 ///
139 /// The output tensor.
140 fn embedding(weights: FloatTensor<B>, indices: IntTensor<B>) -> FloatTensor<B> {
141 let [batch_size, seq_length] = indices.shape().dims();
142 let [_, d_model] = weights.shape().dims();
143
144 let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
145 let output = B::float_select(weights, 0, indices);
146
147 B::float_reshape(output, Shape::new([batch_size, seq_length, d_model]))
148 }
149
150 /// Embedding backward operation.
151 ///
152 /// # Arguments
153 ///
154 /// * `weights` - The embedding weights.
155 /// * `output_grad` - The output gradient.
156 /// * `indices` - The indices tensor.
157 ///
158 /// # Returns
159 ///
160 /// The gradient.
161 fn embedding_backward(
162 weights: FloatTensor<B>,
163 output_grad: FloatTensor<B>,
164 indices: IntTensor<B>,
165 ) -> FloatTensor<B> {
166 let [batch_size, seq_length] = indices.shape().dims();
167 let [n_embeddings, d_model] = weights.shape().dims();
168 let device = weights.device();
169 let dtype = output_grad.dtype();
170
171 let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
172 let output_grad =
173 B::float_reshape(output_grad, Shape::new([batch_size * seq_length, d_model]));
174 let grad = B::float_zeros(Shape::new([n_embeddings, d_model]), &device, dtype.into());
175
176 B::float_select_add(grad, 0, indices, output_grad)
177 }
178
179 /// Linear transformation.
180 ///
181 /// # Shapes
182 ///
183 /// x: `[..., d_input]`,
184 /// weight: `[d_input, d_output]`,
185 /// bias: `[d_output]`,
186 fn linear(
187 x: FloatTensor<B>,
188 weight: FloatTensor<B>,
189 bias: Option<FloatTensor<B>>,
190 ) -> FloatTensor<B> {
191 linear::linear::<B>(x, weight, bias)
192 }
193 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `x`.
194 fn linear_x_backward(weight: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
195 linear::linear_x_backward::<B>(weight, output_grad)
196 }
197 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `weight`.
198 fn linear_weight_backward(x: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
199 linear::linear_weight_backward::<B>(x, output_grad)
200 }
201 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `bias`.
202 fn linear_bias_backward(output_grad: FloatTensor<B>) -> FloatTensor<B> {
203 linear::linear_bias_backward::<B>(output_grad)
204 }
205
206 /// One dimensional convolution.
207 ///
208 /// # Shapes
209 ///
210 /// x: `[batch_size, channels_in, length]`,
211 /// weight: `[channels_out, channels_in, kernel_size]`,
212 /// bias: `[channels_out]`,
213 fn conv1d(
214 x: FloatTensor<B>,
215 weight: FloatTensor<B>,
216 bias: Option<FloatTensor<B>>,
217 options: ConvOptions<1>,
218 ) -> FloatTensor<B> {
219 conv::conv1d_from_conv2d::<B>(x, weight, bias, options)
220 }
221 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `x`.
222 fn conv1d_x_backward(
223 x: FloatTensor<B>,
224 weight: FloatTensor<B>,
225 output_grad: FloatTensor<B>,
226 options: ConvOptions<1>,
227 ) -> FloatTensor<B> {
228 conv::conv1d_x_backward::<B>(x, weight, output_grad, options)
229 }
230 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `weight`.
231 fn conv1d_weight_backward(
232 x: FloatTensor<B>,
233 weight: FloatTensor<B>,
234 output_grad: FloatTensor<B>,
235 options: ConvOptions<1>,
236 ) -> FloatTensor<B> {
237 conv::conv1d_weight_backward::<B>(x, weight, output_grad, options)
238 }
239 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `bias`.
240 fn conv1d_bias_backward(
241 x: FloatTensor<B>,
242 bias: FloatTensor<B>,
243 output_grad: FloatTensor<B>,
244 ) -> FloatTensor<B> {
245 conv::conv1d_bias_backward::<B>(x, bias, output_grad)
246 }
247 /// Two dimensional convolution.
248 ///
249 /// # Shapes
250 ///
251 /// x: `[batch_size, channels_in, height, width]`,
252 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
253 /// bias: `[channels_out]`,
254 fn conv2d(
255 x: FloatTensor<B>,
256 weight: FloatTensor<B>,
257 bias: Option<FloatTensor<B>>,
258 options: ConvOptions<2>,
259 ) -> FloatTensor<B>;
260 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `x`.
261 fn conv2d_x_backward(
262 x: FloatTensor<B>,
263 weight: FloatTensor<B>,
264 output_grad: FloatTensor<B>,
265 options: ConvOptions<2>,
266 ) -> FloatTensor<B> {
267 conv::conv2d_x_backward::<B>(x, weight, output_grad, options)
268 }
269 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `weight`.
270 fn conv2d_weight_backward(
271 x: FloatTensor<B>,
272 weight: FloatTensor<B>,
273 output_grad: FloatTensor<B>,
274 options: ConvOptions<2>,
275 ) -> FloatTensor<B> {
276 conv::conv2d_weight_backward::<B>(x, weight, output_grad, options)
277 }
278 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `bias`.
279 fn conv2d_bias_backward(
280 x: FloatTensor<B>,
281 bias: FloatTensor<B>,
282 output_grad: FloatTensor<B>,
283 ) -> FloatTensor<B> {
284 conv::conv2d_bias_backward::<B>(x, bias, output_grad)
285 }
286
287 /// Two dimensional deformable convolution.
288 ///
289 /// # Shapes
290 ///
291 /// x: `[batch_size, channels_in, height, width]`,
292 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
293 /// bias: `[channels_out]`,
294 fn deform_conv2d(
295 x: FloatTensor<B>,
296 offset: FloatTensor<B>,
297 weight: FloatTensor<B>,
298 mask: Option<FloatTensor<B>>,
299 bias: Option<FloatTensor<B>>,
300 options: DeformConvOptions<2>,
301 ) -> FloatTensor<B>;
302 /// Backward pass for the [deform_conv2d](ModuleOps::deform_conv2d) operation.
303 fn deform_conv2d_backward(
304 x: FloatTensor<B>,
305 offset: FloatTensor<B>,
306 weight: FloatTensor<B>,
307 mask: Option<FloatTensor<B>>,
308 bias: Option<FloatTensor<B>>,
309 output_grad: FloatTensor<B>,
310 options: DeformConvOptions<2>,
311 ) -> DeformConv2dBackward<B>;
312
313 /// Three dimensional convolution.
314 ///
315 /// # Shapes
316 ///
317 /// x: `[batch_size, channels_in, depth, height, width]`,
318 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3]`,
319 /// bias: `[channels_out]`,
320 fn conv3d(
321 x: FloatTensor<B>,
322 weight: FloatTensor<B>,
323 bias: Option<FloatTensor<B>>,
324 options: ConvOptions<3>,
325 ) -> FloatTensor<B>;
326 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `x`.
327 fn conv3d_x_backward(
328 x: FloatTensor<B>,
329 weight: FloatTensor<B>,
330 output_grad: FloatTensor<B>,
331 options: ConvOptions<3>,
332 ) -> FloatTensor<B> {
333 conv::conv3d_x_backward::<B>(x, weight, output_grad, options)
334 }
335 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `weight`.
336 fn conv3d_weight_backward(
337 x: FloatTensor<B>,
338 weight: FloatTensor<B>,
339 output_grad: FloatTensor<B>,
340 options: ConvOptions<3>,
341 ) -> FloatTensor<B> {
342 conv::conv3d_weight_backward::<B>(x, weight, output_grad, options)
343 }
344 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `bias`.
345 fn conv3d_bias_backward(
346 x: FloatTensor<B>,
347 bias: FloatTensor<B>,
348 output_grad: FloatTensor<B>,
349 ) -> FloatTensor<B> {
350 conv::conv3d_bias_backward::<B>(x, bias, output_grad)
351 }
352 /// One dimensional transposed convolution.
353 ///
354 /// # Shapes
355 ///
356 /// x: `[batch_size, channels_in, length]`,
357 /// weight: `[channels_in, channels_out, length]`,
358 /// bias: `[channels_out]`,
359 fn conv_transpose1d(
360 x: FloatTensor<B>,
361 weight: FloatTensor<B>,
362 bias: Option<FloatTensor<B>>,
363 options: ConvTransposeOptions<1>,
364 ) -> FloatTensor<B> {
365 conv::conv_transpose1d_from_conv_transpose2d::<B>(x, weight, bias, options)
366 }
367 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `x`.
368 fn conv_transpose1d_x_backward(
369 weight: FloatTensor<B>,
370 output_grad: FloatTensor<B>,
371 options: ConvTransposeOptions<1>,
372 ) -> FloatTensor<B> {
373 conv::conv_transpose1d_x_backward::<B>(weight, output_grad, options)
374 }
375 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `weight`.
376 fn conv_transpose1d_weight_backward(
377 x: FloatTensor<B>,
378 weight: FloatTensor<B>,
379 output_grad: FloatTensor<B>,
380 options: ConvTransposeOptions<1>,
381 ) -> FloatTensor<B> {
382 conv::conv_transpose1d_weight_backward::<B>(x, weight, output_grad, options)
383 }
384 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `bias`.
385 fn conv_transpose1d_bias_backward(
386 x: FloatTensor<B>,
387 bias: FloatTensor<B>,
388 output_grad: FloatTensor<B>,
389 ) -> FloatTensor<B> {
390 conv::conv_transpose1d_bias_backward::<B>(x, bias, output_grad)
391 }
392
393 /// Two dimensional transposed convolution.
394 ///
395 /// # Shapes
396 ///
397 /// x: `[batch_size, channels_in, height, width]`,
398 /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2]`,
399 /// bias: `[channels_out]`,
400 fn conv_transpose2d(
401 x: FloatTensor<B>,
402 weight: FloatTensor<B>,
403 bias: Option<FloatTensor<B>>,
404 options: ConvTransposeOptions<2>,
405 ) -> FloatTensor<B>;
406 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `x`.
407 fn conv_transpose2d_x_backward(
408 weight: FloatTensor<B>,
409 output_grad: FloatTensor<B>,
410 options: ConvTransposeOptions<2>,
411 ) -> FloatTensor<B> {
412 conv::conv_transpose2d_x_backward::<B>(weight, output_grad, options)
413 }
414 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `weight`.
415 fn conv_transpose2d_weight_backward(
416 x: FloatTensor<B>,
417 weight: FloatTensor<B>,
418 output_grad: FloatTensor<B>,
419 options: ConvTransposeOptions<2>,
420 ) -> FloatTensor<B> {
421 conv::conv_transpose2d_weight_backward::<B>(x, weight, output_grad, options)
422 }
423 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `bias`.
424 fn conv_transpose2d_bias_backward(
425 x: FloatTensor<B>,
426 bias: FloatTensor<B>,
427 output_grad: FloatTensor<B>,
428 ) -> FloatTensor<B> {
429 conv::conv_transpose2d_bias_backward::<B>(x, bias, output_grad)
430 }
431
432 /// Three dimensional transposed convolution.
433 ///
434 /// # Shapes
435 ///
436 /// x: `[batch_size, channels_in, height, width]`,
437 /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3]`,
438 /// bias: `[channels_out]`,
439 fn conv_transpose3d(
440 x: FloatTensor<B>,
441 weight: FloatTensor<B>,
442 bias: Option<FloatTensor<B>>,
443 options: ConvTransposeOptions<3>,
444 ) -> FloatTensor<B>;
445 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `x`.
446 fn conv_transpose3d_x_backward(
447 weight: FloatTensor<B>,
448 output_grad: FloatTensor<B>,
449 options: ConvTransposeOptions<3>,
450 ) -> FloatTensor<B> {
451 conv::conv_transpose3d_x_backward::<B>(weight, output_grad, options)
452 }
453 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `weight`.
454 fn conv_transpose3d_weight_backward(
455 x: FloatTensor<B>,
456 weight: FloatTensor<B>,
457 output_grad: FloatTensor<B>,
458 options: ConvTransposeOptions<3>,
459 ) -> FloatTensor<B> {
460 conv::conv_transpose3d_weight_backward::<B>(x, weight, output_grad, options)
461 }
462 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `bias`.
463 fn conv_transpose3d_bias_backward(
464 x: FloatTensor<B>,
465 bias: FloatTensor<B>,
466 output_grad: FloatTensor<B>,
467 ) -> FloatTensor<B> {
468 conv::conv_transpose3d_bias_backward::<B>(x, bias, output_grad)
469 }
470
471 /// Four-dimensional unfolding.
472 ///
473 /// # Shapes
474 ///
475 /// * x: ``[batch_size, channels_in, height, width]``,
476 /// * returns: ``[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]``,
477 fn unfold4d(
478 x: FloatTensor<B>,
479 kernel_size: [usize; 2],
480 options: UnfoldOptions,
481 ) -> FloatTensor<B> {
482 if options.padding == [0, 0] && options.dilation == [1, 1] {
483 let blocks = B::float_unfold(x, 2, kernel_size[0], options.stride[0]);
484 let blocks = B::float_unfold(blocks, 3, kernel_size[1], options.stride[1]);
485
486 // batch, channels, h_blocks, w_blocks, h_kern, w_kern
487
488 let blocks = B::float_permute(blocks, &[0, 1, 4, 5, 2, 3]);
489 let shape = blocks.shape();
490
491 // batch, channels, h_kern, w_kern, h_blocks, w_blocks
492
493 B::float_reshape(
494 blocks,
495 [
496 shape[0],
497 shape[1] * shape[2] * shape[3],
498 shape[4] * shape[5],
499 ]
500 .into(),
501 )
502 } else {
503 unfold4d_using_conv2d::<B>(x, kernel_size, options)
504 }
505 }
506
507 /// Four dimensional fold (`col2im`), the adjoint of [unfold4d](ModuleOps::unfold4d).
508 ///
509 /// Composes [conv_transpose2d](ModuleOps::conv_transpose2d) with the same one-hot weight
510 /// [unfold4d](ModuleOps::unfold4d) uses, so backends inherit a correct (and differentiable)
511 /// implementation for free and may override it with a custom one.
512 ///
513 /// # Shapes
514 ///
515 /// x: `[batch_size, channels * kernel_size_0 * kernel_size_1, num_blocks]`,
516 /// output: `[batch_size, channels, output_size_0, output_size_1]`
517 fn fold4d(
518 x: FloatTensor<B>,
519 output_size: [usize; 2],
520 kernel_size: [usize; 2],
521 options: UnfoldOptions,
522 ) -> FloatTensor<B> {
523 let [batch_size, channels_col, num_blocks] = x.shape().dims();
524 let [kernel_height, kernel_width] = kernel_size;
525 let [output_height, output_width] = output_size;
526 let [stride_height, stride_width] = options.stride;
527 let [padding_height, padding_width] = options.padding;
528 let [dilation_height, dilation_width] = options.dilation;
529
530 let kernel_elems = kernel_height * kernel_width;
531 assert_eq!(
532 channels_col % kernel_elems,
533 0,
534 "fold4d: input channels ({channels_col}) must be divisible by the kernel size product ({kernel_elems})"
535 );
536 let channels = channels_col / kernel_elems;
537
538 // Number of sliding blocks along each spatial dimension (the unfold output grid).
539 let blocks_height =
540 (output_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1)
541 / stride_height
542 + 1;
543 let blocks_width =
544 (output_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1)
545 / stride_width
546 + 1;
547 assert_eq!(
548 num_blocks,
549 blocks_height * blocks_width,
550 "fold4d: number of blocks ({num_blocks}) does not match the expected grid ({blocks_height} x {blocks_width}) for the given output size and options"
551 );
552
553 // The fold weight is identical to the one `unfold4d` builds for its `conv2d` — fold is its adjoint.
554 let weight = create_unfolding_weight::<B>(channels, kernel_size, &x.device(), x.dtype());
555
556 // Reshape the columns into the spatial grid of blocks, then scatter-add them back.
557 let x = B::float_reshape(
558 x,
559 Shape::new([batch_size, channels_col, blocks_height, blocks_width]),
560 );
561
562 // `padding_out` recovers the exact requested output size (always `< stride`).
563 let padding_out = [
564 (output_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1)
565 % stride_height,
566 (output_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1)
567 % stride_width,
568 ];
569
570 B::conv_transpose2d(
571 x,
572 weight,
573 None,
574 ConvTransposeOptions::new(
575 options.stride,
576 options.padding,
577 padding_out,
578 options.dilation,
579 1,
580 ),
581 )
582 }
583
584 /// One dimensional avg pooling.
585 ///
586 /// # Shapes
587 ///
588 /// x: [batch_size, channels, length],
589 fn avg_pool1d(
590 x: FloatTensor<B>,
591 kernel_size: usize,
592 stride: usize,
593 padding: usize,
594 count_include_pad: bool,
595 ceil_mode: bool,
596 ) -> FloatTensor<B> {
597 pool::avg_pool1d_from_2d::<B>(
598 x,
599 kernel_size,
600 stride,
601 padding,
602 count_include_pad,
603 ceil_mode,
604 )
605 }
606 /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation.
607 fn avg_pool1d_backward(
608 x: FloatTensor<B>,
609 grad: FloatTensor<B>,
610 kernel_size: usize,
611 stride: usize,
612 padding: usize,
613 count_include_pad: bool,
614 ceil_mode: bool,
615 ) -> FloatTensor<B> {
616 pool::avg_pool1d_backward_from_2d::<B>(
617 x,
618 grad,
619 kernel_size,
620 stride,
621 padding,
622 count_include_pad,
623 ceil_mode,
624 )
625 }
626 /// Two dimensional avg pooling.
627 ///
628 /// # Shapes
629 ///
630 /// x: [batch_size, channels, height, width],
631 fn avg_pool2d(
632 x: FloatTensor<B>,
633 kernel_size: [usize; 2],
634 stride: [usize; 2],
635 padding: [usize; 2],
636 count_include_pad: bool,
637 ceil_mode: bool,
638 ) -> FloatTensor<B>;
639 /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
640 fn avg_pool2d_backward(
641 x: FloatTensor<B>,
642 grad: FloatTensor<B>,
643 kernel_size: [usize; 2],
644 stride: [usize; 2],
645 padding: [usize; 2],
646 count_include_pad: bool,
647 ceil_mode: bool,
648 ) -> FloatTensor<B>;
649 /// Two dimensional adaptive avg pooling.
650 ///
651 /// # Shapes
652 ///
653 /// x: [batch_size, channels, height, width],
654 fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
655 /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
656 fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
657 /// Three dimensional adaptive avg pooling.
658 ///
659 /// # Shapes
660 ///
661 /// x: [batch_size, channels, depth, height, width],
662 fn adaptive_avg_pool3d(x: FloatTensor<B>, output_size: [usize; 3]) -> FloatTensor<B>;
663 /// Backward pass for the [adaptive avg pooling 3d](ModuleOps::adaptive_avg_pool3d) operation.
664 fn adaptive_avg_pool3d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
665 /// One dimensional adaptive avg pooling.
666 ///
667 /// # Shapes
668 ///
669 /// x: [batch_size, channels, length],
670 fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
671 pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
672 }
673 /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
674 fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
675 pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
676 }
677 /// One dimensional max pooling.
678 ///
679 /// # Shapes
680 ///
681 /// x: [batch_size, channels, length],
682 fn max_pool1d(
683 x: FloatTensor<B>,
684 kernel_size: usize,
685 stride: usize,
686 padding: usize,
687 dilation: usize,
688 ceil_mode: bool,
689 ) -> FloatTensor<B> {
690 pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
691 }
692
693 /// One dimensional max pooling with indices.
694 ///
695 /// # Shapes
696 ///
697 /// x: [batch_size, channels, height, width],
698 fn max_pool1d_with_indices(
699 x: FloatTensor<B>,
700 kernel_size: usize,
701 stride: usize,
702 padding: usize,
703 dilation: usize,
704 ceil_mode: bool,
705 indices_dtype: IntDType,
706 ) -> MaxPool1dWithIndices<B> {
707 pool::max_pool1d_with_indices_from_2d::<B>(
708 x,
709 kernel_size,
710 stride,
711 padding,
712 dilation,
713 ceil_mode,
714 indices_dtype,
715 )
716 }
717 /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
718 #[allow(clippy::too_many_arguments)]
719 fn max_pool1d_with_indices_backward(
720 x: FloatTensor<B>,
721 kernel_size: usize,
722 stride: usize,
723 padding: usize,
724 dilation: usize,
725 ceil_mode: bool,
726 output_grad: FloatTensor<B>,
727 indices: IntTensor<B>,
728 ) -> MaxPool1dBackward<B> {
729 pool::max_pool1d_with_indices_backward_from_2d::<B>(
730 x,
731 kernel_size,
732 stride,
733 padding,
734 dilation,
735 ceil_mode,
736 output_grad,
737 indices,
738 )
739 }
740
741 /// Two dimensional max pooling.
742 ///
743 /// # Shapes
744 ///
745 /// x: [batch_size, channels, height, width],
746 fn max_pool2d(
747 x: FloatTensor<B>,
748 kernel_size: [usize; 2],
749 stride: [usize; 2],
750 padding: [usize; 2],
751 dilation: [usize; 2],
752 ceil_mode: bool,
753 ) -> FloatTensor<B>;
754
755 /// Two dimensional max pooling with indices.
756 ///
757 /// # Shapes
758 ///
759 /// x: [batch_size, channels, height, width],
760 fn max_pool2d_with_indices(
761 x: FloatTensor<B>,
762 kernel_size: [usize; 2],
763 stride: [usize; 2],
764 padding: [usize; 2],
765 dilation: [usize; 2],
766 ceil_mode: bool,
767 indices_dtype: IntDType,
768 ) -> MaxPool2dWithIndices<B>;
769 /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
770 #[allow(clippy::too_many_arguments)]
771 fn max_pool2d_with_indices_backward(
772 x: FloatTensor<B>,
773 kernel_size: [usize; 2],
774 stride: [usize; 2],
775 padding: [usize; 2],
776 dilation: [usize; 2],
777 ceil_mode: bool,
778 output_grad: FloatTensor<B>,
779 indices: IntTensor<B>,
780 ) -> MaxPool2dBackward<B>;
781
782 /// Down/up samples the input.
783 ///
784 /// # Shapes
785 ///
786 /// x: `[batch_size, channels, height, width]`,
787 fn interpolate(
788 x: FloatTensor<B>,
789 output_size: [usize; 2],
790 options: InterpolateOptions,
791 ) -> FloatTensor<B>;
792
793 /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
794 fn interpolate_backward(
795 x: FloatTensor<B>,
796 grad: FloatTensor<B>,
797 output_size: [usize; 2],
798 options: InterpolateOptions,
799 ) -> FloatTensor<B>;
800
801 /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
802 /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking,
803 /// additive bias, causal masking, and softcap to the attention scores.
804 ///
805 /// # Arguments
806 /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
807 /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
808 /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
809 /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
810 /// where `true` indicates positions to mask (i.e. set to -inf before softmax).
811 /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
812 /// added to the attention scores before softmax (e.g. ALiBi, relative position biases).
813 /// - `options`: Additional attention options (custom scale, softcap, causal masking).
814 ///
815 /// # Returns
816 /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
817 /// representing the attended context per head.
818 ///
819 /// # Note
820 /// This implementation does not support dropout and is intended for inference or
821 /// use cases where dropout is not needed.
822 fn attention(
823 query: FloatTensor<B>,
824 key: FloatTensor<B>,
825 value: FloatTensor<B>,
826 mask: Option<BoolTensor<B>>,
827 attn_bias: Option<FloatTensor<B>>,
828 options: AttentionModuleOptions,
829 ) -> FloatTensor<B>;
830
831 /// Applies Layer Normalization over the last dimension of the input tensor.
832 ///
833 /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
834 /// (biased) `var` are reduced over the last axis.
835 ///
836 /// # Arguments
837 ///
838 /// * `tensor` - Input tensor of shape `[..., d_model]`.
839 /// * `gamma` - Scale tensor of shape `[d_model]`.
840 /// * `beta` - Optional bias tensor of shape `[d_model]`.
841 /// * `epsilon` - Numerical stability term added to the variance before the square root.
842 ///
843 /// # Returns
844 ///
845 /// A tensor with the same shape as `tensor`.
846 fn layer_norm(
847 tensor: FloatTensor<B>,
848 gamma: FloatTensor<B>,
849 beta: Option<FloatTensor<B>>,
850 epsilon: f64,
851 ) -> FloatTensor<B> {
852 let shape = tensor.shape();
853 let rank = shape.num_dims();
854 let last_dim = rank - 1;
855 let d_model = shape[last_dim];
856
857 let mean = B::float_mean_dim(tensor.clone(), last_dim);
858 let centered = B::float_sub(tensor, mean);
859 let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim);
860 let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into()));
861 let normalized = B::float_div(centered, denom);
862
863 let broadcast_dims: alloc::vec::Vec<usize> = (0..rank)
864 .map(|i| if i == last_dim { d_model } else { 1 })
865 .collect();
866 let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone()));
867 let scaled = B::float_mul(normalized, gamma_b);
868
869 match beta {
870 Some(beta) => {
871 let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims));
872 B::float_add(scaled, beta_b)
873 }
874 None => scaled,
875 }
876 }
877
878 /// Computes the Connectionist Temporal Classification (CTC) loss.
879 ///
880 /// Sums over all valid alignments between the input and target sequences
881 /// using the forward (alpha) algorithm.
882 ///
883 /// # Arguments
884 ///
885 /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
886 /// * `targets` - Target label indices of shape `[N, S]`
887 /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
888 /// * `target_lengths` - Actual target lengths per batch element `[N]`
889 /// * `blank` - Index of the blank label
890 ///
891 /// # Returns
892 ///
893 /// Per-sample loss of shape `[N]`
894 fn ctc_loss(
895 log_probs: FloatTensor<B>,
896 targets: IntTensor<B>,
897 input_lengths: IntTensor<B>,
898 target_lengths: IntTensor<B>,
899 blank: usize,
900 ) -> FloatTensor<B> {
901 ctc::ctc_loss_default::<B>(log_probs, targets, input_lengths, target_lengths, blank)
902 }
903
904 /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward)
905 /// natively.
906 ///
907 /// Autodiff queries this flag to decide between two paths:
908 /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and
909 /// [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly.
910 /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff
911 /// then differentiates through the decomposed tensor ops.
912 ///
913 /// Backends that override `ctc_loss_backward` must also override this to
914 /// return `true`.
915 fn has_ctc_loss_backward() -> bool {
916 false
917 }
918
919 /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`.
920 ///
921 /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward)
922 /// returns `true`. Backends without a native implementation should leave
923 /// both methods at their defaults; the gradient is computed automatically by
924 /// autodiff against the decomposed [ctc::ctc_loss_default] forward.
925 ///
926 /// # Arguments
927 ///
928 /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
929 /// * `targets` - Target label indices of shape `[N, S]`
930 /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
931 /// * `target_lengths` - Actual target lengths per batch element `[N]`
932 /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]`
933 /// * `blank` - Index of the blank label
934 ///
935 /// # Returns
936 ///
937 /// Gradient w.r.t. `log_probs` of shape `[T, N, C]`
938 fn ctc_loss_backward(
939 _log_probs: FloatTensor<B>,
940 _targets: IntTensor<B>,
941 _input_lengths: IntTensor<B>,
942 _target_lengths: IntTensor<B>,
943 _grad_loss: FloatTensor<B>,
944 _blank: usize,
945 ) -> FloatTensor<B> {
946 unreachable!(
947 "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false"
948 )
949 }
950
951 /// Real-valued FFT with optional size parameter.
952 ///
953 /// When `n` is `None`, the signal must be a power of two along `dim`, and the output has
954 /// `signal_len / 2 + 1` frequency bins.
955 ///
956 /// When `n` is `Some(size)`, `size` must also be a power of two. The signal is truncated
957 /// or zero-padded to `size` and the output has `size / 2 + 1` frequency bins. Non-power-
958 /// of-two sizes are currently rejected at the public API boundary; true arbitrary-`n` DFT
959 /// support (Bluestein's algorithm) is tracked as a follow-up.
960 ///
961 /// Returns two tensors: the real part and the imaginary part.
962 fn rfft(
963 signal: FloatTensor<B>,
964 dim: usize,
965 n: Option<usize>,
966 ) -> (FloatTensor<B>, FloatTensor<B>);
967
968 /// Inverse real-valued FFT with optional output size.
969 ///
970 /// When `n` is `None`, the reconstructed signal length `2 * (spectrum_size - 1)` must be
971 /// a power of two.
972 ///
973 /// When `n` is `Some(size)`, `size` must also be a power of two. Output has exactly
974 /// `size` samples.
975 fn irfft(
976 spectrum_re: FloatTensor<B>,
977 spectrum_im: FloatTensor<B>,
978 dim: usize,
979 n: Option<usize>,
980 ) -> FloatTensor<B>;
981}