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, AvgPoolOptions, ConvOptions, ConvTransposeOptions, DeformConvOptions,
8    GridSampleOptions, GridSamplePaddingMode, InterpolateMode, InterpolateOptions, MaxPoolOptions,
9    PadMode, 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 channel mapping
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        // Fold is the adjoint of unfold and uses the same one-hot channel mapping. Grouping by
688        // channel removes the otherwise unused cross-channel weights.
689        let groups = channels.max(1);
690        let weight =
691            create_unfolding_weight::<B>(channels, kernel_size, groups, &x.device(), x.dtype());
692
693        // Reshape the columns into the spatial grid of blocks, then scatter-add them back.
694        let x = B::float_reshape(
695            x,
696            Shape::new([batch_size, channels_col, blocks_height, blocks_width]),
697        );
698
699        // `padding_out` recovers the exact requested output size (always `< stride`).
700        let padding_out = [
701            (output_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1)
702                % stride_height,
703            (output_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1)
704                % stride_width,
705        ];
706
707        B::conv_transpose2d(
708            x,
709            weight,
710            None,
711            ConvTransposeOptions::new(
712                options.stride,
713                options.padding,
714                padding_out,
715                options.dilation,
716                groups,
717            ),
718        )
719    }
720
721    /// One dimensional avg pooling.
722    ///
723    /// # Shapes
724    ///
725    /// x: [batch_size, channels, length],
726    fn avg_pool1d(
727        x: FloatTensor<B>,
728        kernel_size: usize,
729        stride: usize,
730        padding: usize,
731        count_include_pad: bool,
732        ceil_mode: bool,
733    ) -> FloatTensor<B> {
734        pool::avg_pool1d_from_2d::<B>(
735            x,
736            kernel_size,
737            stride,
738            padding,
739            count_include_pad,
740            ceil_mode,
741        )
742    }
743    /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation.
744    fn avg_pool1d_backward(
745        x: FloatTensor<B>,
746        grad: FloatTensor<B>,
747        kernel_size: usize,
748        stride: usize,
749        padding: usize,
750        count_include_pad: bool,
751        ceil_mode: bool,
752    ) -> FloatTensor<B> {
753        pool::avg_pool1d_backward_from_2d::<B>(
754            x,
755            grad,
756            kernel_size,
757            stride,
758            padding,
759            count_include_pad,
760            ceil_mode,
761        )
762    }
763    /// Two dimensional avg pooling.
764    ///
765    /// # Shapes
766    ///
767    /// x: [batch_size, channels, height, width],
768    fn avg_pool2d(
769        x: FloatTensor<B>,
770        kernel_size: [usize; 2],
771        stride: [usize; 2],
772        padding: [usize; 2],
773        count_include_pad: bool,
774        ceil_mode: bool,
775    ) -> FloatTensor<B>;
776    /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
777    fn avg_pool2d_backward(
778        x: FloatTensor<B>,
779        grad: FloatTensor<B>,
780        kernel_size: [usize; 2],
781        stride: [usize; 2],
782        padding: [usize; 2],
783        count_include_pad: bool,
784        ceil_mode: bool,
785    ) -> FloatTensor<B>;
786    /// Two dimensional adaptive avg pooling.
787    ///
788    /// # Shapes
789    ///
790    /// x: [batch_size, channels, height, width],
791    fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
792    /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
793    fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
794    /// Three dimensional adaptive avg pooling.
795    ///
796    /// # Shapes
797    ///
798    /// x: [batch_size, channels, depth, height, width],
799    fn adaptive_avg_pool3d(x: FloatTensor<B>, output_size: [usize; 3]) -> FloatTensor<B>;
800    /// Backward pass for the [adaptive avg pooling 3d](ModuleOps::adaptive_avg_pool3d) operation.
801    fn adaptive_avg_pool3d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
802    /// One dimensional adaptive avg pooling.
803    ///
804    /// # Shapes
805    ///
806    /// x: [batch_size, channels, length],
807    fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
808        pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
809    }
810    /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
811    fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
812        pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
813    }
814    /// One dimensional max pooling.
815    ///
816    /// # Shapes
817    ///
818    /// x: [batch_size, channels, length],
819    fn max_pool1d(
820        x: FloatTensor<B>,
821        kernel_size: usize,
822        stride: usize,
823        padding: usize,
824        dilation: usize,
825        ceil_mode: bool,
826    ) -> FloatTensor<B> {
827        pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
828    }
829
830    /// One dimensional max pooling with indices.
831    ///
832    /// # Shapes
833    ///
834    /// x: [batch_size, channels, height, width],
835    fn max_pool1d_with_indices(
836        x: FloatTensor<B>,
837        kernel_size: usize,
838        stride: usize,
839        padding: usize,
840        dilation: usize,
841        ceil_mode: bool,
842        indices_dtype: IntDType,
843    ) -> MaxPool1dWithIndices<B> {
844        pool::max_pool1d_with_indices_from_2d::<B>(
845            x,
846            kernel_size,
847            stride,
848            padding,
849            dilation,
850            ceil_mode,
851            indices_dtype,
852        )
853    }
854    /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
855    #[allow(clippy::too_many_arguments)]
856    fn max_pool1d_with_indices_backward(
857        x: FloatTensor<B>,
858        kernel_size: usize,
859        stride: usize,
860        padding: usize,
861        dilation: usize,
862        ceil_mode: bool,
863        output_grad: FloatTensor<B>,
864        indices: IntTensor<B>,
865    ) -> MaxPool1dBackward<B> {
866        pool::max_pool1d_with_indices_backward_from_2d::<B>(
867            x,
868            kernel_size,
869            stride,
870            padding,
871            dilation,
872            ceil_mode,
873            output_grad,
874            indices,
875        )
876    }
877
878    /// Two dimensional max pooling.
879    ///
880    /// # Shapes
881    ///
882    /// x: [batch_size, channels, height, width],
883    fn max_pool2d(
884        x: FloatTensor<B>,
885        kernel_size: [usize; 2],
886        stride: [usize; 2],
887        padding: [usize; 2],
888        dilation: [usize; 2],
889        ceil_mode: bool,
890    ) -> FloatTensor<B>;
891
892    /// Two dimensional max pooling with indices.
893    ///
894    /// # Shapes
895    ///
896    /// x: [batch_size, channels, height, width],
897    fn max_pool2d_with_indices(
898        x: FloatTensor<B>,
899        kernel_size: [usize; 2],
900        stride: [usize; 2],
901        padding: [usize; 2],
902        dilation: [usize; 2],
903        ceil_mode: bool,
904        indices_dtype: IntDType,
905    ) -> MaxPool2dWithIndices<B>;
906    /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
907    #[allow(clippy::too_many_arguments)]
908    fn max_pool2d_with_indices_backward(
909        x: FloatTensor<B>,
910        kernel_size: [usize; 2],
911        stride: [usize; 2],
912        padding: [usize; 2],
913        dilation: [usize; 2],
914        ceil_mode: bool,
915        output_grad: FloatTensor<B>,
916        indices: IntTensor<B>,
917    ) -> MaxPool2dBackward<B>;
918
919    /// Down/up samples the input.
920    ///
921    /// # Shapes
922    ///
923    /// x: `[batch_size, channels, height, width]`,
924    fn interpolate(
925        x: FloatTensor<B>,
926        output_size: [usize; 2],
927        options: InterpolateOptions,
928    ) -> FloatTensor<B>;
929
930    /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
931    fn interpolate_backward(
932        x: FloatTensor<B>,
933        grad: FloatTensor<B>,
934        output_size: [usize; 2],
935        options: InterpolateOptions,
936    ) -> FloatTensor<B>;
937
938    /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
939    /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking,
940    /// additive bias, causal masking, and softcap to the attention scores.
941    ///
942    /// # Arguments
943    /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
944    /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
945    /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
946    /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
947    ///   where `true` indicates positions to mask (i.e. set to -inf before softmax).
948    /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
949    ///   added to the attention scores before softmax (e.g. ALiBi, relative position biases).
950    /// - `options`: Additional attention options (custom scale, softcap, causal masking).
951    ///
952    /// # Returns
953    /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
954    /// representing the attended context per head.
955    ///
956    /// # Note
957    /// This implementation does not support dropout and is intended for inference or
958    /// use cases where dropout is not needed.
959    fn attention(
960        query: FloatTensor<B>,
961        key: FloatTensor<B>,
962        value: FloatTensor<B>,
963        mask: Option<BoolTensor<B>>,
964        attn_bias: Option<FloatTensor<B>>,
965        options: AttentionModuleOptions,
966    ) -> FloatTensor<B>;
967
968    /// Applies Layer Normalization over the last dimension of the input tensor.
969    ///
970    /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
971    /// (biased) `var` are reduced over the last axis.
972    ///
973    /// # Arguments
974    ///
975    /// * `tensor` - Input tensor of shape `[..., d_model]`.
976    /// * `gamma` - Scale tensor of shape `[d_model]`.
977    /// * `beta` - Optional bias tensor of shape `[d_model]`.
978    /// * `epsilon` - Numerical stability term added to the variance before the square root.
979    ///
980    /// # Returns
981    ///
982    /// A tensor with the same shape as `tensor`.
983    fn layer_norm(
984        tensor: FloatTensor<B>,
985        gamma: FloatTensor<B>,
986        beta: Option<FloatTensor<B>>,
987        epsilon: f64,
988    ) -> FloatTensor<B> {
989        let shape = tensor.shape();
990        let rank = shape.num_dims();
991        let last_dim = rank - 1;
992        let d_model = shape[last_dim];
993
994        let mean = B::float_mean_dim(tensor.clone(), last_dim);
995        let centered = B::float_sub(tensor, mean);
996        let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim);
997        let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into()));
998        let normalized = B::float_div(centered, denom);
999
1000        let broadcast_dims: alloc::vec::Vec<usize> = (0..rank)
1001            .map(|i| if i == last_dim { d_model } else { 1 })
1002            .collect();
1003        let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone()));
1004        let scaled = B::float_mul(normalized, gamma_b);
1005
1006        match beta {
1007            Some(beta) => {
1008                let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims));
1009                B::float_add(scaled, beta_b)
1010            }
1011            None => scaled,
1012        }
1013    }
1014
1015    /// Computes the Connectionist Temporal Classification (CTC) loss.
1016    ///
1017    /// Sums over all valid alignments between the input and target sequences
1018    /// using the forward (alpha) algorithm.
1019    ///
1020    /// # Arguments
1021    ///
1022    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
1023    /// * `targets` - Target label indices of shape `[N, S]`
1024    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
1025    /// * `target_lengths` - Actual target lengths per batch element `[N]`
1026    /// * `blank` - Index of the blank label
1027    ///
1028    /// # Returns
1029    ///
1030    /// Per-sample loss of shape `[N]`
1031    fn ctc_loss(
1032        log_probs: FloatTensor<B>,
1033        targets: IntTensor<B>,
1034        input_lengths: IntTensor<B>,
1035        target_lengths: IntTensor<B>,
1036        blank: usize,
1037    ) -> FloatTensor<B> {
1038        ctc::ctc_loss_default::<B>(log_probs, targets, input_lengths, target_lengths, blank)
1039    }
1040
1041    /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward)
1042    /// natively.
1043    ///
1044    /// Autodiff queries this flag to decide between two paths:
1045    /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and
1046    ///   [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly.
1047    /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff
1048    ///   then differentiates through the decomposed tensor ops.
1049    ///
1050    /// Backends that override `ctc_loss_backward` must also override this to
1051    /// return `true`.
1052    fn has_ctc_loss_backward() -> bool {
1053        false
1054    }
1055
1056    /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`.
1057    ///
1058    /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward)
1059    /// returns `true`. Backends without a native implementation should leave
1060    /// both methods at their defaults; the gradient is computed automatically by
1061    /// autodiff against the decomposed [ctc::ctc_loss_default] forward.
1062    ///
1063    /// # Arguments
1064    ///
1065    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
1066    /// * `targets` - Target label indices of shape `[N, S]`
1067    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
1068    /// * `target_lengths` - Actual target lengths per batch element `[N]`
1069    /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]`
1070    /// * `blank` - Index of the blank label
1071    ///
1072    /// # Returns
1073    ///
1074    /// Gradient w.r.t. `log_probs` of shape `[T, N, C]`
1075    fn ctc_loss_backward(
1076        _log_probs: FloatTensor<B>,
1077        _targets: IntTensor<B>,
1078        _input_lengths: IntTensor<B>,
1079        _target_lengths: IntTensor<B>,
1080        _grad_loss: FloatTensor<B>,
1081        _blank: usize,
1082    ) -> FloatTensor<B> {
1083        unreachable!(
1084            "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false"
1085        )
1086    }
1087}
burn_backend/backend/ops/modules/base.rs

burn_backend/backend/ops/modules/
base.rs
