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

burn_backend/backend/ops/modules/
base.rs
