Skip to main content

burn_backend/backend/ops/modules/
base.rs

1use super::{conv, ctc, linear, pool};
2use crate::ops::unfold::unfold4d_using_conv2d;
3use crate::tensor::{BoolTensor, FloatTensor, IntTensor};
4use crate::{Backend, 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    /// Embedding operation.
101    ///
102    /// # Arguments
103    ///
104    /// * `weights` - The embedding weights.
105    /// * `indices` - The indices tensor.
106    ///
107    /// # Returns
108    ///
109    /// The output tensor.
110    fn embedding(weights: FloatTensor<B>, indices: IntTensor<B>) -> FloatTensor<B> {
111        let [batch_size, seq_length] = indices.shape().dims();
112        let [_, d_model] = weights.shape().dims();
113
114        let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
115        let output = B::float_select(weights, 0, indices);
116
117        B::float_reshape(output, Shape::new([batch_size, seq_length, d_model]))
118    }
119
120    /// Embedding backward operation.
121    ///
122    /// # Arguments
123    ///
124    /// * `weights` - The embedding weights.
125    /// * `output_grad` - The output gradient.
126    /// * `indices` - The indices tensor.
127    ///
128    /// # Returns
129    ///
130    /// The gradient.
131    fn embedding_backward(
132        weights: FloatTensor<B>,
133        output_grad: FloatTensor<B>,
134        indices: IntTensor<B>,
135    ) -> FloatTensor<B> {
136        let [batch_size, seq_length] = indices.shape().dims();
137        let [n_embeddings, d_model] = weights.shape().dims();
138        let device = weights.device();
139        let dtype = output_grad.dtype();
140
141        let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
142        let output_grad =
143            B::float_reshape(output_grad, Shape::new([batch_size * seq_length, d_model]));
144        let grad = B::float_zeros(Shape::new([n_embeddings, d_model]), &device, dtype.into());
145
146        B::float_select_add(grad, 0, indices, output_grad)
147    }
148
149    /// Linear transformation.
150    ///
151    /// # Shapes
152    ///
153    /// x:      `[..., d_input]`,
154    /// weight: `[d_input, d_output]`,
155    /// bias:   `[d_output]`,
156    fn linear(
157        x: FloatTensor<B>,
158        weight: FloatTensor<B>,
159        bias: Option<FloatTensor<B>>,
160    ) -> FloatTensor<B> {
161        linear::linear::<B>(x, weight, bias)
162    }
163    /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `x`.
164    fn linear_x_backward(weight: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
165        linear::linear_x_backward::<B>(weight, output_grad)
166    }
167    /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `weight`.
168    fn linear_weight_backward(x: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
169        linear::linear_weight_backward::<B>(x, output_grad)
170    }
171    /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `bias`.
172    fn linear_bias_backward(output_grad: FloatTensor<B>) -> FloatTensor<B> {
173        linear::linear_bias_backward::<B>(output_grad)
174    }
175
176    /// One dimensional convolution.
177    ///
178    /// # Shapes
179    ///
180    /// x:      `[batch_size, channels_in, length]`,
181    /// weight: `[channels_out, channels_in, kernel_size]`,
182    /// bias:   `[channels_out]`,
183    fn conv1d(
184        x: FloatTensor<B>,
185        weight: FloatTensor<B>,
186        bias: Option<FloatTensor<B>>,
187        options: ConvOptions<1>,
188    ) -> FloatTensor<B> {
189        conv::conv1d_from_conv2d::<B>(x, weight, bias, options)
190    }
191    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `x`.
192    fn conv1d_x_backward(
193        x: FloatTensor<B>,
194        weight: FloatTensor<B>,
195        output_grad: FloatTensor<B>,
196        options: ConvOptions<1>,
197    ) -> FloatTensor<B> {
198        conv::conv1d_x_backward::<B>(x, weight, output_grad, options)
199    }
200    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `weight`.
201    fn conv1d_weight_backward(
202        x: FloatTensor<B>,
203        weight: FloatTensor<B>,
204        output_grad: FloatTensor<B>,
205        options: ConvOptions<1>,
206    ) -> FloatTensor<B> {
207        conv::conv1d_weight_backward::<B>(x, weight, output_grad, options)
208    }
209    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `bias`.
210    fn conv1d_bias_backward(
211        x: FloatTensor<B>,
212        bias: FloatTensor<B>,
213        output_grad: FloatTensor<B>,
214    ) -> FloatTensor<B> {
215        conv::conv1d_bias_backward::<B>(x, bias, output_grad)
216    }
217    /// Two dimensional convolution.
218    ///
219    /// # Shapes
220    ///
221    /// x:      `[batch_size, channels_in, height, width]`,
222    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
223    /// bias:   `[channels_out]`,
224    fn conv2d(
225        x: FloatTensor<B>,
226        weight: FloatTensor<B>,
227        bias: Option<FloatTensor<B>>,
228        options: ConvOptions<2>,
229    ) -> FloatTensor<B>;
230    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `x`.
231    fn conv2d_x_backward(
232        x: FloatTensor<B>,
233        weight: FloatTensor<B>,
234        output_grad: FloatTensor<B>,
235        options: ConvOptions<2>,
236    ) -> FloatTensor<B> {
237        conv::conv2d_x_backward::<B>(x, weight, output_grad, options)
238    }
239    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `weight`.
240    fn conv2d_weight_backward(
241        x: FloatTensor<B>,
242        weight: FloatTensor<B>,
243        output_grad: FloatTensor<B>,
244        options: ConvOptions<2>,
245    ) -> FloatTensor<B> {
246        conv::conv2d_weight_backward::<B>(x, weight, output_grad, options)
247    }
248    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `bias`.
249    fn conv2d_bias_backward(
250        x: FloatTensor<B>,
251        bias: FloatTensor<B>,
252        output_grad: FloatTensor<B>,
253    ) -> FloatTensor<B> {
254        conv::conv2d_bias_backward::<B>(x, bias, output_grad)
255    }
256
257    /// Two dimensional deformable convolution.
258    ///
259    /// # Shapes
260    ///
261    /// x:      `[batch_size, channels_in, height, width]`,
262    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
263    /// bias:   `[channels_out]`,
264    fn deform_conv2d(
265        x: FloatTensor<B>,
266        offset: FloatTensor<B>,
267        weight: FloatTensor<B>,
268        mask: Option<FloatTensor<B>>,
269        bias: Option<FloatTensor<B>>,
270        options: DeformConvOptions<2>,
271    ) -> FloatTensor<B>;
272    /// Backward pass for the [deform_conv2d](ModuleOps::deform_conv2d) operation.
273    fn deform_conv2d_backward(
274        x: FloatTensor<B>,
275        offset: FloatTensor<B>,
276        weight: FloatTensor<B>,
277        mask: Option<FloatTensor<B>>,
278        bias: Option<FloatTensor<B>>,
279        output_grad: FloatTensor<B>,
280        options: DeformConvOptions<2>,
281    ) -> DeformConv2dBackward<B>;
282
283    /// Three dimensional convolution.
284    ///
285    /// # Shapes
286    ///
287    /// x:      `[batch_size, channels_in, depth, height, width]`,
288    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3]`,
289    /// bias:   `[channels_out]`,
290    fn conv3d(
291        x: FloatTensor<B>,
292        weight: FloatTensor<B>,
293        bias: Option<FloatTensor<B>>,
294        options: ConvOptions<3>,
295    ) -> FloatTensor<B>;
296    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `x`.
297    fn conv3d_x_backward(
298        x: FloatTensor<B>,
299        weight: FloatTensor<B>,
300        output_grad: FloatTensor<B>,
301        options: ConvOptions<3>,
302    ) -> FloatTensor<B> {
303        conv::conv3d_x_backward::<B>(x, weight, output_grad, options)
304    }
305    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `weight`.
306    fn conv3d_weight_backward(
307        x: FloatTensor<B>,
308        weight: FloatTensor<B>,
309        output_grad: FloatTensor<B>,
310        options: ConvOptions<3>,
311    ) -> FloatTensor<B> {
312        conv::conv3d_weight_backward::<B>(x, weight, output_grad, options)
313    }
314    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `bias`.
315    fn conv3d_bias_backward(
316        x: FloatTensor<B>,
317        bias: FloatTensor<B>,
318        output_grad: FloatTensor<B>,
319    ) -> FloatTensor<B> {
320        conv::conv3d_bias_backward::<B>(x, bias, output_grad)
321    }
322    /// One dimensional transposed convolution.
323    ///
324    /// # Shapes
325    ///
326    /// x:      `[batch_size, channels_in, length]`,
327    /// weight: `[channels_in, channels_out, length]`,
328    /// bias:   `[channels_out]`,
329    fn conv_transpose1d(
330        x: FloatTensor<B>,
331        weight: FloatTensor<B>,
332        bias: Option<FloatTensor<B>>,
333        options: ConvTransposeOptions<1>,
334    ) -> FloatTensor<B> {
335        conv::conv_transpose1d_from_conv_transpose2d::<B>(x, weight, bias, options)
336    }
337    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `x`.
338    fn conv_transpose1d_x_backward(
339        weight: FloatTensor<B>,
340        output_grad: FloatTensor<B>,
341        options: ConvTransposeOptions<1>,
342    ) -> FloatTensor<B> {
343        conv::conv_transpose1d_x_backward::<B>(weight, output_grad, options)
344    }
345    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `weight`.
346    fn conv_transpose1d_weight_backward(
347        x: FloatTensor<B>,
348        weight: FloatTensor<B>,
349        output_grad: FloatTensor<B>,
350        options: ConvTransposeOptions<1>,
351    ) -> FloatTensor<B> {
352        conv::conv_transpose1d_weight_backward::<B>(x, weight, output_grad, options)
353    }
354    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `bias`.
355    fn conv_transpose1d_bias_backward(
356        x: FloatTensor<B>,
357        bias: FloatTensor<B>,
358        output_grad: FloatTensor<B>,
359    ) -> FloatTensor<B> {
360        conv::conv_transpose1d_bias_backward::<B>(x, bias, output_grad)
361    }
362
363    /// Two dimensional transposed convolution.
364    ///
365    /// # Shapes
366    ///
367    /// x:      `[batch_size, channels_in, height, width]`,
368    /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2]`,
369    /// bias:   `[channels_out]`,
370    fn conv_transpose2d(
371        x: FloatTensor<B>,
372        weight: FloatTensor<B>,
373        bias: Option<FloatTensor<B>>,
374        options: ConvTransposeOptions<2>,
375    ) -> FloatTensor<B>;
376    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `x`.
377    fn conv_transpose2d_x_backward(
378        weight: FloatTensor<B>,
379        output_grad: FloatTensor<B>,
380        options: ConvTransposeOptions<2>,
381    ) -> FloatTensor<B> {
382        conv::conv_transpose2d_x_backward::<B>(weight, output_grad, options)
383    }
384    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `weight`.
385    fn conv_transpose2d_weight_backward(
386        x: FloatTensor<B>,
387        weight: FloatTensor<B>,
388        output_grad: FloatTensor<B>,
389        options: ConvTransposeOptions<2>,
390    ) -> FloatTensor<B> {
391        conv::conv_transpose2d_weight_backward::<B>(x, weight, output_grad, options)
392    }
393    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `bias`.
394    fn conv_transpose2d_bias_backward(
395        x: FloatTensor<B>,
396        bias: FloatTensor<B>,
397        output_grad: FloatTensor<B>,
398    ) -> FloatTensor<B> {
399        conv::conv_transpose2d_bias_backward::<B>(x, bias, output_grad)
400    }
401
402    /// Three dimensional transposed convolution.
403    ///
404    /// # Shapes
405    ///
406    /// x:      `[batch_size, channels_in, height, width]`,
407    /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3]`,
408    /// bias:   `[channels_out]`,
409    fn conv_transpose3d(
410        x: FloatTensor<B>,
411        weight: FloatTensor<B>,
412        bias: Option<FloatTensor<B>>,
413        options: ConvTransposeOptions<3>,
414    ) -> FloatTensor<B>;
415    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `x`.
416    fn conv_transpose3d_x_backward(
417        weight: FloatTensor<B>,
418        output_grad: FloatTensor<B>,
419        options: ConvTransposeOptions<3>,
420    ) -> FloatTensor<B> {
421        conv::conv_transpose3d_x_backward::<B>(weight, output_grad, options)
422    }
423    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `weight`.
424    fn conv_transpose3d_weight_backward(
425        x: FloatTensor<B>,
426        weight: FloatTensor<B>,
427        output_grad: FloatTensor<B>,
428        options: ConvTransposeOptions<3>,
429    ) -> FloatTensor<B> {
430        conv::conv_transpose3d_weight_backward::<B>(x, weight, output_grad, options)
431    }
432    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `bias`.
433    fn conv_transpose3d_bias_backward(
434        x: FloatTensor<B>,
435        bias: FloatTensor<B>,
436        output_grad: FloatTensor<B>,
437    ) -> FloatTensor<B> {
438        conv::conv_transpose3d_bias_backward::<B>(x, bias, output_grad)
439    }
440
441    /// Four-dimensional unfolding.
442    ///
443    /// # Shapes
444    ///
445    /// * x:      ``[batch_size, channels_in, height, width]``,
446    /// * returns: ``[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]``,
447    fn unfold4d(
448        x: FloatTensor<B>,
449        kernel_size: [usize; 2],
450        options: UnfoldOptions,
451    ) -> FloatTensor<B> {
452        if options.padding == [0, 0] && options.dilation == [1, 1] {
453            let blocks = B::float_unfold(x, 2, kernel_size[0], options.stride[0]);
454            let blocks = B::float_unfold(blocks, 3, kernel_size[1], options.stride[1]);
455
456            // batch, channels, h_blocks, w_blocks, h_kern, w_kern
457
458            let blocks = B::float_permute(blocks, &[0, 1, 4, 5, 2, 3]);
459            let shape = blocks.shape();
460
461            // batch, channels, h_kern, w_kern, h_blocks, w_blocks
462
463            B::float_reshape(
464                blocks,
465                [
466                    shape[0],
467                    shape[1] * shape[2] * shape[3],
468                    shape[4] * shape[5],
469                ]
470                .into(),
471            )
472        } else {
473            unfold4d_using_conv2d::<B>(x, kernel_size, options)
474        }
475    }
476
477    /// One dimensional avg pooling.
478    ///
479    /// # Shapes
480    ///
481    /// x: [batch_size, channels, length],
482    fn avg_pool1d(
483        x: FloatTensor<B>,
484        kernel_size: usize,
485        stride: usize,
486        padding: usize,
487        count_include_pad: bool,
488        ceil_mode: bool,
489    ) -> FloatTensor<B> {
490        pool::avg_pool1d_from_2d::<B>(
491            x,
492            kernel_size,
493            stride,
494            padding,
495            count_include_pad,
496            ceil_mode,
497        )
498    }
499    /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation.
500    fn avg_pool1d_backward(
501        x: FloatTensor<B>,
502        grad: FloatTensor<B>,
503        kernel_size: usize,
504        stride: usize,
505        padding: usize,
506        count_include_pad: bool,
507        ceil_mode: bool,
508    ) -> FloatTensor<B> {
509        pool::avg_pool1d_backward_from_2d::<B>(
510            x,
511            grad,
512            kernel_size,
513            stride,
514            padding,
515            count_include_pad,
516            ceil_mode,
517        )
518    }
519    /// Two dimensional avg pooling.
520    ///
521    /// # Shapes
522    ///
523    /// x: [batch_size, channels, height, width],
524    fn avg_pool2d(
525        x: FloatTensor<B>,
526        kernel_size: [usize; 2],
527        stride: [usize; 2],
528        padding: [usize; 2],
529        count_include_pad: bool,
530        ceil_mode: bool,
531    ) -> FloatTensor<B>;
532    /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
533    fn avg_pool2d_backward(
534        x: FloatTensor<B>,
535        grad: FloatTensor<B>,
536        kernel_size: [usize; 2],
537        stride: [usize; 2],
538        padding: [usize; 2],
539        count_include_pad: bool,
540        ceil_mode: bool,
541    ) -> FloatTensor<B>;
542    /// Two dimensional adaptive avg pooling.
543    ///
544    /// # Shapes
545    ///
546    /// x: [batch_size, channels, height, width],
547    fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
548    /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
549    fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
550    /// One dimensional adaptive avg pooling.
551    ///
552    /// # Shapes
553    ///
554    /// x: [batch_size, channels, length],
555    fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
556        pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
557    }
558    /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
559    fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
560        pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
561    }
562    /// One dimensional max pooling.
563    ///
564    /// # Shapes
565    ///
566    /// x: [batch_size, channels, length],
567    fn max_pool1d(
568        x: FloatTensor<B>,
569        kernel_size: usize,
570        stride: usize,
571        padding: usize,
572        dilation: usize,
573        ceil_mode: bool,
574    ) -> FloatTensor<B> {
575        pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
576    }
577
578    /// One dimensional max pooling with indices.
579    ///
580    /// # Shapes
581    ///
582    /// x: [batch_size, channels, height, width],
583    fn max_pool1d_with_indices(
584        x: FloatTensor<B>,
585        kernel_size: usize,
586        stride: usize,
587        padding: usize,
588        dilation: usize,
589        ceil_mode: bool,
590        indices_dtype: IntDType,
591    ) -> MaxPool1dWithIndices<B> {
592        pool::max_pool1d_with_indices_from_2d::<B>(
593            x,
594            kernel_size,
595            stride,
596            padding,
597            dilation,
598            ceil_mode,
599            indices_dtype,
600        )
601    }
602    /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
603    #[allow(clippy::too_many_arguments)]
604    fn max_pool1d_with_indices_backward(
605        x: FloatTensor<B>,
606        kernel_size: usize,
607        stride: usize,
608        padding: usize,
609        dilation: usize,
610        ceil_mode: bool,
611        output_grad: FloatTensor<B>,
612        indices: IntTensor<B>,
613    ) -> MaxPool1dBackward<B> {
614        pool::max_pool1d_with_indices_backward_from_2d::<B>(
615            x,
616            kernel_size,
617            stride,
618            padding,
619            dilation,
620            ceil_mode,
621            output_grad,
622            indices,
623        )
624    }
625
626    /// Two dimensional max pooling.
627    ///
628    /// # Shapes
629    ///
630    /// x: [batch_size, channels, height, width],
631    fn max_pool2d(
632        x: FloatTensor<B>,
633        kernel_size: [usize; 2],
634        stride: [usize; 2],
635        padding: [usize; 2],
636        dilation: [usize; 2],
637        ceil_mode: bool,
638    ) -> FloatTensor<B>;
639
640    /// Two dimensional max pooling with indices.
641    ///
642    /// # Shapes
643    ///
644    /// x: [batch_size, channels, height, width],
645    fn max_pool2d_with_indices(
646        x: FloatTensor<B>,
647        kernel_size: [usize; 2],
648        stride: [usize; 2],
649        padding: [usize; 2],
650        dilation: [usize; 2],
651        ceil_mode: bool,
652        indices_dtype: IntDType,
653    ) -> MaxPool2dWithIndices<B>;
654    /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
655    #[allow(clippy::too_many_arguments)]
656    fn max_pool2d_with_indices_backward(
657        x: FloatTensor<B>,
658        kernel_size: [usize; 2],
659        stride: [usize; 2],
660        padding: [usize; 2],
661        dilation: [usize; 2],
662        ceil_mode: bool,
663        output_grad: FloatTensor<B>,
664        indices: IntTensor<B>,
665    ) -> MaxPool2dBackward<B>;
666
667    /// Down/up samples the input.
668    ///
669    /// # Shapes
670    ///
671    /// x: `[batch_size, channels, height, width]`,
672    fn interpolate(
673        x: FloatTensor<B>,
674        output_size: [usize; 2],
675        options: InterpolateOptions,
676    ) -> FloatTensor<B>;
677
678    /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
679    fn interpolate_backward(
680        x: FloatTensor<B>,
681        grad: FloatTensor<B>,
682        output_size: [usize; 2],
683        options: InterpolateOptions,
684    ) -> FloatTensor<B>;
685
686    /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
687    /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking,
688    /// additive bias, causal masking, and softcap to the attention scores.
689    ///
690    /// # Arguments
691    /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
692    /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
693    /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
694    /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
695    ///   where `true` indicates positions to mask (i.e. set to -inf before softmax).
696    /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
697    ///   added to the attention scores before softmax (e.g. ALiBi, relative position biases).
698    /// - `options`: Additional attention options (custom scale, softcap, causal masking).
699    ///
700    /// # Returns
701    /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
702    /// representing the attended context per head.
703    ///
704    /// # Note
705    /// This implementation does not support dropout and is intended for inference or
706    /// use cases where dropout is not needed.
707    fn attention(
708        query: FloatTensor<B>,
709        key: FloatTensor<B>,
710        value: FloatTensor<B>,
711        mask: Option<BoolTensor<B>>,
712        attn_bias: Option<FloatTensor<B>>,
713        options: AttentionModuleOptions,
714    ) -> FloatTensor<B>;
715
716    /// Applies Layer Normalization over the last dimension of the input tensor.
717    ///
718    /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
719    /// (biased) `var` are reduced over the last axis.
720    ///
721    /// # Arguments
722    ///
723    /// * `tensor` - Input tensor of shape `[..., d_model]`.
724    /// * `gamma` - Scale tensor of shape `[d_model]`.
725    /// * `beta` - Optional bias tensor of shape `[d_model]`.
726    /// * `epsilon` - Numerical stability term added to the variance before the square root.
727    ///
728    /// # Returns
729    ///
730    /// A tensor with the same shape as `tensor`.
731    fn layer_norm(
732        tensor: FloatTensor<B>,
733        gamma: FloatTensor<B>,
734        beta: Option<FloatTensor<B>>,
735        epsilon: f64,
736    ) -> FloatTensor<B> {
737        let shape = tensor.shape();
738        let rank = shape.num_dims();
739        let last_dim = rank - 1;
740        let d_model = shape[last_dim];
741
742        let mean = B::float_mean_dim(tensor.clone(), last_dim);
743        let centered = B::float_sub(tensor, mean);
744        let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim);
745        let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into()));
746        let normalized = B::float_div(centered, denom);
747
748        let broadcast_dims: alloc::vec::Vec<usize> = (0..rank)
749            .map(|i| if i == last_dim { d_model } else { 1 })
750            .collect();
751        let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone()));
752        let scaled = B::float_mul(normalized, gamma_b);
753
754        match beta {
755            Some(beta) => {
756                let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims));
757                B::float_add(scaled, beta_b)
758            }
759            None => scaled,
760        }
761    }
762
763    /// Computes the Connectionist Temporal Classification (CTC) loss.
764    ///
765    /// Sums over all valid alignments between the input and target sequences
766    /// using the forward (alpha) algorithm.
767    ///
768    /// # Arguments
769    ///
770    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
771    /// * `targets` - Target label indices of shape `[N, S]`
772    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
773    /// * `target_lengths` - Actual target lengths per batch element `[N]`
774    /// * `blank` - Index of the blank label
775    ///
776    /// # Returns
777    ///
778    /// Per-sample loss of shape `[N]`
779    fn ctc_loss(
780        log_probs: FloatTensor<B>,
781        targets: IntTensor<B>,
782        input_lengths: IntTensor<B>,
783        target_lengths: IntTensor<B>,
784        blank: usize,
785    ) -> FloatTensor<B> {
786        ctc::ctc_loss_default::<B>(log_probs, targets, input_lengths, target_lengths, blank)
787    }
788
789    /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward)
790    /// natively.
791    ///
792    /// Autodiff queries this flag to decide between two paths:
793    /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and
794    ///   [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly.
795    /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff
796    ///   then differentiates through the decomposed tensor ops.
797    ///
798    /// Backends that override `ctc_loss_backward` must also override this to
799    /// return `true`.
800    fn has_ctc_loss_backward() -> bool {
801        false
802    }
803
804    /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`.
805    ///
806    /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward)
807    /// returns `true`. Backends without a native implementation should leave
808    /// both methods at their defaults; the gradient is computed automatically by
809    /// autodiff against the decomposed [ctc::ctc_loss_default] forward.
810    ///
811    /// # Arguments
812    ///
813    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
814    /// * `targets` - Target label indices of shape `[N, S]`
815    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
816    /// * `target_lengths` - Actual target lengths per batch element `[N]`
817    /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]`
818    /// * `blank` - Index of the blank label
819    ///
820    /// # Returns
821    ///
822    /// Gradient w.r.t. `log_probs` of shape `[T, N, C]`
823    fn ctc_loss_backward(
824        _log_probs: FloatTensor<B>,
825        _targets: IntTensor<B>,
826        _input_lengths: IntTensor<B>,
827        _target_lengths: IntTensor<B>,
828        _grad_loss: FloatTensor<B>,
829        _blank: usize,
830    ) -> FloatTensor<B> {
831        unreachable!(
832            "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false"
833        )
834    }
835
836    /// Real-valued FFT with optional size parameter.
837    ///
838    /// When `n` is `None`, the signal must be a power of two along `dim`, and the output has
839    /// `signal_len / 2 + 1` frequency bins.
840    ///
841    /// When `n` is `Some(size)`, `size` must also be a power of two. The signal is truncated
842    /// or zero-padded to `size` and the output has `size / 2 + 1` frequency bins. Non-power-
843    /// of-two sizes are currently rejected at the public API boundary; true arbitrary-`n` DFT
844    /// support (Bluestein's algorithm) is tracked as a follow-up.
845    ///
846    /// Returns two tensors: the real part and the imaginary part.
847    fn rfft(
848        signal: FloatTensor<B>,
849        dim: usize,
850        n: Option<usize>,
851    ) -> (FloatTensor<B>, FloatTensor<B>);
852
853    /// Inverse real-valued FFT with optional output size.
854    ///
855    /// When `n` is `None`, the reconstructed signal length `2 * (spectrum_size - 1)` must be
856    /// a power of two.
857    ///
858    /// When `n` is `Some(size)`, `size` must also be a power of two. Output has exactly
859    /// `size` samples.
860    fn irfft(
861        spectrum_re: FloatTensor<B>,
862        spectrum_im: FloatTensor<B>,
863        dim: usize,
864        n: Option<usize>,
865    ) -> FloatTensor<B>;
866}