burn_backend/backend/ops/modules/
base.rs

1use super::{conv, pool};
2use crate::ops::attention;
3use crate::ops::unfold::unfold4d_using_conv2d;
4use crate::tensor::{BoolTensor, FloatTensor, IntTensor};
5use crate::{Backend, TensorMetadata};
6use burn_std::Shape;
7use core::num::NonZeroUsize;
8
9/// Gradient computed during the backward pass for each tensor used by [conv2d](ModuleOps::conv2d).
10#[derive(new)]
11pub struct Conv2dBackward<B: Backend> {
12    /// Gradient.
13    pub x_grad: FloatTensor<B>,
14
15    /// Weights gradient.
16    pub weights_grad: FloatTensor<B>,
17
18    /// Bias gradient.
19    pub bias_grad: Option<FloatTensor<B>>,
20}
21
22/// Gradient computed during the backward pass for each tensor used by [deform_conv2d](ModuleOps::deform_conv2d).
23#[derive(new)]
24pub struct DeformConv2dBackward<B: Backend> {
25    /// Gradient.
26    pub x_grad: FloatTensor<B>,
27
28    /// Offset gradient.
29    pub offset_grad: FloatTensor<B>,
30
31    /// Weights gradient.
32    pub weight_grad: FloatTensor<B>,
33
34    /// Mask gradient.
35    pub mask_grad: Option<FloatTensor<B>>,
36
37    /// Bias gradient.
38    pub bias_grad: Option<FloatTensor<B>>,
39}
40
41/// Gradient computed during the backward pass for each tensor used by [conv3d](ModuleOps::conv3d).
42#[derive(new)]
43pub struct Conv3dBackward<B: Backend> {
44    /// Gradient.
45    pub x_grad: FloatTensor<B>,
46
47    /// Weights gradient.
48    pub weights_grad: FloatTensor<B>,
49
50    /// Bias gradient.
51    pub bias_grad: Option<FloatTensor<B>>,
52}
53
54/// Gradient computed during the backward pass for each tensor used by [max_pool1d](ModuleOps::max_pool1d).
55#[derive(new)]
56pub struct MaxPool1dBackward<B: Backend> {
57    /// Gradient.
58    pub x_grad: FloatTensor<B>,
59}
60
61/// Results from [max_pool1d](ModuleOps::max_pool1d_with_indices).
62#[derive(new)]
63pub struct MaxPool1dWithIndices<B: Backend> {
64    /// The output tensor.
65    pub output: FloatTensor<B>,
66
67    /// The indices tensor.
68    pub indices: IntTensor<B>,
69}
70
71/// Gradient computed during the backward pass for each tensor used by [max_pool2d](ModuleOps::max_pool2d).
72#[derive(new)]
73pub struct MaxPool2dBackward<B: Backend> {
74    /// Gradient.
75    pub x_grad: FloatTensor<B>,
76}
77
78/// Results from [max_pool2d](ModuleOps::max_pool2d_with_indices).
79#[derive(new)]
80pub struct MaxPool2dWithIndices<B: Backend> {
81    /// The output tensor.
82    pub output: FloatTensor<B>,
83
84    /// The indices tensor.
85    pub indices: IntTensor<B>,
86}
87
88/// Check that the parameter value is non-zero.
89// NOTE: for now we keep usize but we could refactor the parameters to hold `NonZeroUsize`.
90pub(crate) fn check_nonzero(value: usize, msg: &str) -> usize {
91    NonZeroUsize::new(value).expect(msg);
92    value
93}
94
95/// Convolution options.
96#[derive(Debug, Clone, Hash, PartialEq, Eq)]
97pub struct ConvOptions<const N: usize> {
98    /// Stride (non-zero).
99    pub stride: [usize; N],
100
101    /// Padding.
102    pub padding: [usize; N],
103
104    /// Dilation (non-zero).
105    pub dilation: [usize; N],
106
107    /// Groups (non-zero).
108    pub groups: usize,
109}
110
111impl<const N: usize> ConvOptions<N> {
112    /// Constructs a new `ConvOptions`.
113    pub fn new(
114        stride: [usize; N],
115        padding: [usize; N],
116        dilation: [usize; N],
117        groups: usize,
118    ) -> Self {
119        Self {
120            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
121            padding,
122            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
123            groups: check_nonzero(groups, "groups must be non-zero"),
124        }
125    }
126}
127
128/// Convolution options.
129#[derive(Debug, Clone, Hash, PartialEq, Eq)]
130pub struct DeformConvOptions<const N: usize> {
131    /// Stride (non-zero).
132    pub stride: [usize; N],
133
134    /// Padding.
135    pub padding: [usize; N],
136
137    /// Dilation (non-zero).
138    pub dilation: [usize; N],
139
140    /// Weight Groups (non-zero).
141    pub weight_groups: usize,
142
143    /// Offset Groups (non-zero).
144    pub offset_groups: usize,
145}
146
147impl<const N: usize> DeformConvOptions<N> {
148    /// Constructs a new `DeformConvOptions`.
149    pub fn new(
150        stride: [usize; N],
151        padding: [usize; N],
152        dilation: [usize; N],
153        weight_groups: usize,
154        offset_groups: usize,
155    ) -> Self {
156        Self {
157            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
158            padding,
159            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
160            weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"),
161            offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"),
162        }
163    }
164}
165
166/// Transposed convolution options.
167#[derive(Debug, Clone, Hash, PartialEq, Eq)]
168pub struct ConvTransposeOptions<const N: usize> {
169    /// Stride (non-zero).
170    pub stride: [usize; N],
171
172    /// Padding.
173    pub padding: [usize; N],
174
175    /// Padding out.
176    pub padding_out: [usize; N],
177
178    /// Dilation (non-zero).
179    pub dilation: [usize; N],
180
181    /// Groups (non-zero).
182    pub groups: usize,
183}
184
185impl<const N: usize> ConvTransposeOptions<N> {
186    /// Constructs a new `ConvTransposeOptions`.
187    pub fn new(
188        stride: [usize; N],
189        padding: [usize; N],
190        padding_out: [usize; N],
191        dilation: [usize; N],
192        groups: usize,
193    ) -> Self {
194        Self {
195            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
196            padding,
197            padding_out,
198            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
199            groups: check_nonzero(groups, "groups must be non-zero"),
200        }
201    }
202}
203
204/// Unfold operation options.
205#[derive(Debug, Clone)]
206pub struct UnfoldOptions {
207    /// The number of positions to slide over the input tensor in each dimension.
208    /// A stride of `[1, 1]` will slide the kernel one pixel at a time.
209    pub stride: [usize; 2],
210
211    /// The number of zero-padding pixels added to each side of the input tensor in each dimension.
212    pub padding: [usize; 2],
213
214    /// The spacing between the blocks (patches) in the original input tensor.
215    pub dilation: [usize; 2],
216}
217
218impl UnfoldOptions {
219    /// Constructs a new `UnfoldOptions`.
220    pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self {
221        Self {
222            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
223            padding,
224            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
225        }
226    }
227}
228
229/// Algorithm used for upsampling.
230#[derive(new, Debug, Clone, serde::Deserialize, serde::Serialize)]
231pub enum InterpolateMode {
232    /// Nearest-neighbor interpolation.
233    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
234    Nearest,
235
236    /// Bilinear interpolation.
237    /// <https://en.wikipedia.org/wiki/Bilinear_interpolation>
238    Bilinear,
239
240    /// Bicubic interpolation.
241    /// <https://en.wikipedia.org/wiki/Bicubic_interpolation>
242    Bicubic,
243}
244
245/// Interpolation options.
246#[derive(new, Debug, Clone)]
247pub struct InterpolateOptions {
248    /// Algorithm used for upsampling.
249    pub mode: InterpolateMode,
250}
251
252/// Padding mode for grid sampling when coordinates are out of bounds.
253///
254/// Matches PyTorch's `padding_mode` parameter in `grid_sample`.
255#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)]
256pub enum GridSamplePaddingMode {
257    /// Fill with zeros for out-of-bounds coordinates.
258    #[default]
259    Zeros,
260    /// Clamp coordinates to the border (use nearest edge value).
261    Border,
262    /// Reflect coordinates at the boundary.
263    Reflection,
264}
265
266/// Options for grid sampling operations.
267#[derive(Debug, Clone)]
268pub struct GridSampleOptions {
269    /// Interpolation mode (bilinear, nearest, or bicubic).
270    pub mode: InterpolateMode,
271    /// Padding mode for out-of-bounds coordinates.
272    pub padding_mode: GridSamplePaddingMode,
273    /// If `true`, grid values of -1 and 1 correspond to the corner pixels.
274    /// If `false`, they correspond to the corner points of the corner pixels
275    /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates).
276    pub align_corners: bool,
277}
278
279impl Default for GridSampleOptions {
280    fn default() -> Self {
281        Self {
282            mode: InterpolateMode::Bilinear,
283            padding_mode: GridSamplePaddingMode::Zeros,
284            align_corners: false,
285        }
286    }
287}
288
289impl GridSampleOptions {
290    /// Create new grid sample options with the given interpolation mode.
291    ///
292    /// Uses default values for padding_mode (Zeros) and align_corners (false).
293    pub fn new(mode: InterpolateMode) -> Self {
294        Self {
295            mode,
296            ..Default::default()
297        }
298    }
299
300    /// Set the padding mode.
301    pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self {
302        self.padding_mode = padding_mode;
303        self
304    }
305
306    /// Set align_corners.
307    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
308        self.align_corners = align_corners;
309        self
310    }
311}
312
313/// Padding mode for tensor pad operations.
314///
315/// Defines how values are filled when padding a tensor beyond its original boundaries.
316///
317/// **Note**: Currently, padding is only supported on the last two dimensions of a tensor
318/// (typically height and width for image data in NCHW format).
319///
320/// # Modes
321///
322/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0)
323/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size)
324/// - [`Edge`](PadMode::Edge): Replicate boundary values
325#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)]
326pub enum PadMode {
327    /// Fill padded regions with a constant value.
328    ///
329    /// # Example
330    /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0:
331    /// Result: `[0, 0, 1, 2, 3]`
332    Constant(f32),
333
334    /// Reflect values at the boundary, excluding the edge value.
335    ///
336    /// Padding must be less than the dimension size (i.e., `padding < dim_size`).
337    ///
338    /// # Example
339    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
340    /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0)
341    Reflect,
342
343    /// Replicate the edge values.
344    ///
345    /// # Example
346    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
347    /// Result: `[1, 1, 1, 2, 3, 4]`
348    Edge,
349}
350
351impl Default for PadMode {
352    fn default() -> Self {
353        PadMode::Constant(0.0)
354    }
355}
356
357/// Gradient computed during the backward pass for each tensor used by [interpolate](ModuleOps::interpolate).
358#[derive(new)]
359pub struct InterpolateBackward<B: Backend> {
360    /// Gradient.
361    pub x_grad: FloatTensor<B>,
362}
363
364/// Module operations trait.
365pub trait ModuleOps<B: Backend> {
366    /// Embedding operation.
367    ///
368    /// # Arguments
369    ///
370    /// * `weights` - The embedding weights.
371    /// * `indices` - The indices tensor.
372    ///
373    /// # Returns
374    ///
375    /// The output tensor.
376    fn embedding(weights: FloatTensor<B>, indices: IntTensor<B>) -> FloatTensor<B> {
377        let [batch_size, seq_length] = indices.shape().dims();
378        let [_, d_model] = weights.shape().dims();
379
380        let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
381        let output = B::float_select(weights, 0, indices);
382
383        B::float_reshape(output, Shape::new([batch_size, seq_length, d_model]))
384    }
385
386    /// Embedding backward operation.
387    ///
388    /// # Arguments
389    ///
390    /// * `weights` - The embedding weights.
391    /// * `output_grad` - The output gradient.
392    /// * `indices` - The indices tensor.
393    ///
394    /// # Returns
395    ///
396    /// The gradient.
397    fn embedding_backward(
398        weights: FloatTensor<B>,
399        output_grad: FloatTensor<B>,
400        indices: IntTensor<B>,
401    ) -> FloatTensor<B> {
402        let [batch_size, seq_length] = indices.shape().dims();
403        let [n_embeddings, d_model] = weights.shape().dims();
404        let device = B::float_device(&weights);
405        let dtype = output_grad.dtype();
406
407        let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
408        let output_grad =
409            B::float_reshape(output_grad, Shape::new([batch_size * seq_length, d_model]));
410        let grad = B::float_zeros(Shape::new([n_embeddings, d_model]), &device, dtype.into());
411
412        B::float_select_add(grad, 0, indices, output_grad)
413    }
414    /// One dimensional convolution.
415    ///
416    /// # Shapes
417    ///
418    /// x:      `[batch_size, channels_in, length]`,
419    /// weight: `[channels_out, channels_in, kernel_size]`,
420    /// bias:   `[channels_out]`,
421    fn conv1d(
422        x: FloatTensor<B>,
423        weight: FloatTensor<B>,
424        bias: Option<FloatTensor<B>>,
425        options: ConvOptions<1>,
426    ) -> FloatTensor<B> {
427        conv::conv1d_from_conv2d::<B>(x, weight, bias, options)
428    }
429    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `x`.
430    fn conv1d_x_backward(
431        x: FloatTensor<B>,
432        weight: FloatTensor<B>,
433        output_grad: FloatTensor<B>,
434        options: ConvOptions<1>,
435    ) -> FloatTensor<B> {
436        conv::conv1d_x_backward::<B>(x, weight, output_grad, options)
437    }
438    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `weight`.
439    fn conv1d_weight_backward(
440        x: FloatTensor<B>,
441        weight: FloatTensor<B>,
442        output_grad: FloatTensor<B>,
443        options: ConvOptions<1>,
444    ) -> FloatTensor<B> {
445        conv::conv1d_weight_backward::<B>(x, weight, output_grad, options)
446    }
447    /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `bias`.
448    fn conv1d_bias_backward(
449        x: FloatTensor<B>,
450        bias: FloatTensor<B>,
451        output_grad: FloatTensor<B>,
452    ) -> FloatTensor<B> {
453        conv::conv1d_bias_backward::<B>(x, bias, output_grad)
454    }
455    /// Two dimensional convolution.
456    ///
457    /// # Shapes
458    ///
459    /// x:      `[batch_size, channels_in, height, width]`,
460    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
461    /// bias:   `[channels_out]`,
462    fn conv2d(
463        x: FloatTensor<B>,
464        weight: FloatTensor<B>,
465        bias: Option<FloatTensor<B>>,
466        options: ConvOptions<2>,
467    ) -> FloatTensor<B>;
468    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `x`.
469    fn conv2d_x_backward(
470        x: FloatTensor<B>,
471        weight: FloatTensor<B>,
472        output_grad: FloatTensor<B>,
473        options: ConvOptions<2>,
474    ) -> FloatTensor<B> {
475        conv::conv2d_x_backward::<B>(x, weight, output_grad, options)
476    }
477    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `weight`.
478    fn conv2d_weight_backward(
479        x: FloatTensor<B>,
480        weight: FloatTensor<B>,
481        output_grad: FloatTensor<B>,
482        options: ConvOptions<2>,
483    ) -> FloatTensor<B> {
484        conv::conv2d_weight_backward::<B>(x, weight, output_grad, options)
485    }
486    /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `bias`.
487    fn conv2d_bias_backward(
488        x: FloatTensor<B>,
489        weight: FloatTensor<B>,
490        bias: FloatTensor<B>,
491        output_grad: FloatTensor<B>,
492    ) -> FloatTensor<B> {
493        conv::conv2d_bias_backward::<B>(x, weight, bias, output_grad)
494    }
495
496    /// Two dimensional deformable convolution.
497    ///
498    /// # Shapes
499    ///
500    /// x:      `[batch_size, channels_in, height, width]`,
501    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
502    /// bias:   `[channels_out]`,
503    fn deform_conv2d(
504        x: FloatTensor<B>,
505        offset: FloatTensor<B>,
506        weight: FloatTensor<B>,
507        mask: Option<FloatTensor<B>>,
508        bias: Option<FloatTensor<B>>,
509        options: DeformConvOptions<2>,
510    ) -> FloatTensor<B>;
511    /// Backward pass for the [deform_conv2d](ModuleOps::deform_conv2d) operation.
512    fn deform_conv2d_backward(
513        x: FloatTensor<B>,
514        offset: FloatTensor<B>,
515        weight: FloatTensor<B>,
516        mask: Option<FloatTensor<B>>,
517        bias: Option<FloatTensor<B>>,
518        output_grad: FloatTensor<B>,
519        options: DeformConvOptions<2>,
520    ) -> DeformConv2dBackward<B>;
521
522    /// Three dimensional convolution.
523    ///
524    /// # Shapes
525    ///
526    /// x:      `[batch_size, channels_in, depth, height, width]`,
527    /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3]`,
528    /// bias:   `[channels_out]`,
529    fn conv3d(
530        x: FloatTensor<B>,
531        weight: FloatTensor<B>,
532        bias: Option<FloatTensor<B>>,
533        options: ConvOptions<3>,
534    ) -> FloatTensor<B>;
535    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `x`.
536    fn conv3d_x_backward(
537        x: FloatTensor<B>,
538        weight: FloatTensor<B>,
539        output_grad: FloatTensor<B>,
540        options: ConvOptions<3>,
541    ) -> FloatTensor<B> {
542        conv::conv3d_x_backward::<B>(x, weight, output_grad, options)
543    }
544    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `weight`.
545    fn conv3d_weight_backward(
546        x: FloatTensor<B>,
547        weight: FloatTensor<B>,
548        output_grad: FloatTensor<B>,
549        options: ConvOptions<3>,
550    ) -> FloatTensor<B> {
551        conv::conv3d_weight_backward::<B>(x, weight, output_grad, options)
552    }
553    /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `bias`.
554    fn conv3d_bias_backward(
555        x: FloatTensor<B>,
556        weight: FloatTensor<B>,
557        bias: FloatTensor<B>,
558        output_grad: FloatTensor<B>,
559    ) -> FloatTensor<B> {
560        conv::conv3d_bias_backward::<B>(x, weight, bias, output_grad)
561    }
562    /// One dimensional transposed convolution.
563    ///
564    /// # Shapes
565    ///
566    /// x:      `[batch_size, channels_in, length]`,
567    /// weight: `[channels_in, channels_out, length]`,
568    /// bias:   `[channels_out]`,
569    fn conv_transpose1d(
570        x: FloatTensor<B>,
571        weight: FloatTensor<B>,
572        bias: Option<FloatTensor<B>>,
573        options: ConvTransposeOptions<1>,
574    ) -> FloatTensor<B> {
575        conv::conv_transpose1d_from_conv_transpose2d::<B>(x, weight, bias, options)
576    }
577    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `x`.
578    fn conv_transpose1d_x_backward(
579        weight: FloatTensor<B>,
580        output_grad: FloatTensor<B>,
581        options: ConvTransposeOptions<1>,
582    ) -> FloatTensor<B> {
583        conv::conv_transpose1d_x_backward::<B>(weight, output_grad, options)
584    }
585    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `weight`.
586    fn conv_transpose1d_weight_backward(
587        x: FloatTensor<B>,
588        weight: FloatTensor<B>,
589        output_grad: FloatTensor<B>,
590        options: ConvTransposeOptions<1>,
591    ) -> FloatTensor<B> {
592        conv::conv_transpose1d_weight_backward::<B>(x, weight, output_grad, options)
593    }
594    /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `bias`.
595    fn conv_transpose1d_bias_backward(
596        x: FloatTensor<B>,
597        bias: FloatTensor<B>,
598        output_grad: FloatTensor<B>,
599    ) -> FloatTensor<B> {
600        conv::conv_transpose1d_bias_backward::<B>(x, bias, output_grad)
601    }
602
603    /// Two dimensional transposed convolution.
604    ///
605    /// # Shapes
606    ///
607    /// x:      `[batch_size, channels_in, height, width]`,
608    /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2]`,
609    /// bias:   `[channels_out]`,
610    fn conv_transpose2d(
611        x: FloatTensor<B>,
612        weight: FloatTensor<B>,
613        bias: Option<FloatTensor<B>>,
614        options: ConvTransposeOptions<2>,
615    ) -> FloatTensor<B>;
616    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `x`.
617    fn conv_transpose2d_x_backward(
618        weight: FloatTensor<B>,
619        output_grad: FloatTensor<B>,
620        options: ConvTransposeOptions<2>,
621    ) -> FloatTensor<B> {
622        conv::conv_transpose2d_x_backward::<B>(weight, output_grad, options)
623    }
624    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `weight`.
625    fn conv_transpose2d_weight_backward(
626        x: FloatTensor<B>,
627        weight: FloatTensor<B>,
628        output_grad: FloatTensor<B>,
629        options: ConvTransposeOptions<2>,
630    ) -> FloatTensor<B> {
631        conv::conv_transpose2d_weight_backward::<B>(x, weight, output_grad, options)
632    }
633    /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `bias`.
634    fn conv_transpose2d_bias_backward(
635        x: FloatTensor<B>,
636        bias: FloatTensor<B>,
637        output_grad: FloatTensor<B>,
638    ) -> FloatTensor<B> {
639        conv::conv_transpose2d_bias_backward::<B>(x, bias, output_grad)
640    }
641
642    /// Three dimensional transposed convolution.
643    ///
644    /// # Shapes
645    ///
646    /// x:      `[batch_size, channels_in, height, width]`,
647    /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3]`,
648    /// bias:   `[channels_out]`,
649    fn conv_transpose3d(
650        x: FloatTensor<B>,
651        weight: FloatTensor<B>,
652        bias: Option<FloatTensor<B>>,
653        options: ConvTransposeOptions<3>,
654    ) -> FloatTensor<B>;
655    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `x`.
656    fn conv_transpose3d_x_backward(
657        weight: FloatTensor<B>,
658        output_grad: FloatTensor<B>,
659        options: ConvTransposeOptions<3>,
660    ) -> FloatTensor<B> {
661        conv::conv_transpose3d_x_backward::<B>(weight, output_grad, options)
662    }
663    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `weight`.
664    fn conv_transpose3d_weight_backward(
665        x: FloatTensor<B>,
666        weight: FloatTensor<B>,
667        output_grad: FloatTensor<B>,
668        options: ConvTransposeOptions<3>,
669    ) -> FloatTensor<B> {
670        conv::conv_transpose3d_weight_backward::<B>(x, weight, output_grad, options)
671    }
672    /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `bias`.
673    fn conv_transpose3d_bias_backward(
674        x: FloatTensor<B>,
675        bias: FloatTensor<B>,
676        output_grad: FloatTensor<B>,
677    ) -> FloatTensor<B> {
678        conv::conv_transpose3d_bias_backward::<B>(x, bias, output_grad)
679    }
680
681    /// Four-dimensional unfolding.
682    ///
683    /// # Shapes
684    ///
685    /// * x:      ``[batch_size, channels_in, height, width]``,
686    /// * returns: ``[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]``,
687    fn unfold4d(
688        x: FloatTensor<B>,
689        kernel_size: [usize; 2],
690        options: UnfoldOptions,
691    ) -> FloatTensor<B> {
692        if options.padding == [0, 0] && options.dilation == [1, 1] {
693            let blocks = B::float_unfold(x, 2, kernel_size[0], options.stride[0]);
694            let blocks = B::float_unfold(blocks, 3, kernel_size[1], options.stride[1]);
695
696            // batch, channels, h_blocks, w_blocks, h_kern, w_kern
697
698            let blocks = B::float_permute(blocks, &[0, 1, 4, 5, 2, 3]);
699            let shape = &blocks.shape().dims;
700
701            // batch, channels, h_kern, w_kern, h_blocks, w_blocks
702
703            B::float_reshape(
704                blocks,
705                [
706                    shape[0],
707                    shape[1] * shape[2] * shape[3],
708                    shape[4] * shape[5],
709                ]
710                .into(),
711            )
712        } else {
713            unfold4d_using_conv2d::<B>(x, kernel_size, options)
714        }
715    }
716
717    /// One dimensional avg pooling.
718    ///
719    /// # Shapes
720    ///
721    /// x: [batch_size, channels, length],
722    fn avg_pool1d(
723        x: FloatTensor<B>,
724        kernel_size: usize,
725        stride: usize,
726        padding: usize,
727        count_include_pad: bool,
728        ceil_mode: bool,
729    ) -> FloatTensor<B> {
730        pool::avg_pool1d_from_2d::<B>(
731            x,
732            kernel_size,
733            stride,
734            padding,
735            count_include_pad,
736            ceil_mode,
737        )
738    }
739    /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation.
740    fn avg_pool1d_backward(
741        x: FloatTensor<B>,
742        grad: FloatTensor<B>,
743        kernel_size: usize,
744        stride: usize,
745        padding: usize,
746        count_include_pad: bool,
747        ceil_mode: bool,
748    ) -> FloatTensor<B> {
749        pool::avg_pool1d_backward_from_2d::<B>(
750            x,
751            grad,
752            kernel_size,
753            stride,
754            padding,
755            count_include_pad,
756            ceil_mode,
757        )
758    }
759    /// Two dimensional avg pooling.
760    ///
761    /// # Shapes
762    ///
763    /// x: [batch_size, channels, height, width],
764    fn avg_pool2d(
765        x: FloatTensor<B>,
766        kernel_size: [usize; 2],
767        stride: [usize; 2],
768        padding: [usize; 2],
769        count_include_pad: bool,
770        ceil_mode: bool,
771    ) -> FloatTensor<B>;
772    /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
773    fn avg_pool2d_backward(
774        x: FloatTensor<B>,
775        grad: FloatTensor<B>,
776        kernel_size: [usize; 2],
777        stride: [usize; 2],
778        padding: [usize; 2],
779        count_include_pad: bool,
780        ceil_mode: bool,
781    ) -> FloatTensor<B>;
782    /// Two dimensional adaptive avg pooling.
783    ///
784    /// # Shapes
785    ///
786    /// x: [batch_size, channels, height, width],
787    fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
788    /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
789    fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
790    /// One dimensional adaptive avg pooling.
791    ///
792    /// # Shapes
793    ///
794    /// x: [batch_size, channels, length],
795    fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
796        pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
797    }
798    /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
799    fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
800        pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
801    }
802    /// One dimensional max pooling.
803    ///
804    /// # Shapes
805    ///
806    /// x: [batch_size, channels, length],
807    fn max_pool1d(
808        x: FloatTensor<B>,
809        kernel_size: usize,
810        stride: usize,
811        padding: usize,
812        dilation: usize,
813        ceil_mode: bool,
814    ) -> FloatTensor<B> {
815        pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
816    }
817
818    /// One dimensional max pooling with indices.
819    ///
820    /// # Shapes
821    ///
822    /// x: [batch_size, channels, height, width],
823    fn max_pool1d_with_indices(
824        x: FloatTensor<B>,
825        kernel_size: usize,
826        stride: usize,
827        padding: usize,
828        dilation: usize,
829        ceil_mode: bool,
830    ) -> MaxPool1dWithIndices<B> {
831        pool::max_pool1d_with_indices_from_2d::<B>(
832            x,
833            kernel_size,
834            stride,
835            padding,
836            dilation,
837            ceil_mode,
838        )
839    }
840    /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
841    #[allow(clippy::too_many_arguments)]
842    fn max_pool1d_with_indices_backward(
843        x: FloatTensor<B>,
844        kernel_size: usize,
845        stride: usize,
846        padding: usize,
847        dilation: usize,
848        ceil_mode: bool,
849        output_grad: FloatTensor<B>,
850        indices: IntTensor<B>,
851    ) -> MaxPool1dBackward<B> {
852        pool::max_pool1d_with_indices_backward_from_2d::<B>(
853            x,
854            kernel_size,
855            stride,
856            padding,
857            dilation,
858            ceil_mode,
859            output_grad,
860            indices,
861        )
862    }
863
864    /// Two dimensional max pooling.
865    ///
866    /// # Shapes
867    ///
868    /// x: [batch_size, channels, height, width],
869    fn max_pool2d(
870        x: FloatTensor<B>,
871        kernel_size: [usize; 2],
872        stride: [usize; 2],
873        padding: [usize; 2],
874        dilation: [usize; 2],
875        ceil_mode: bool,
876    ) -> FloatTensor<B>;
877
878    /// Two dimensional max pooling with indices.
879    ///
880    /// # Shapes
881    ///
882    /// x: [batch_size, channels, height, width],
883    fn max_pool2d_with_indices(
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    ) -> MaxPool2dWithIndices<B>;
891    /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
892    #[allow(clippy::too_many_arguments)]
893    fn max_pool2d_with_indices_backward(
894        x: FloatTensor<B>,
895        kernel_size: [usize; 2],
896        stride: [usize; 2],
897        padding: [usize; 2],
898        dilation: [usize; 2],
899        ceil_mode: bool,
900        output_grad: FloatTensor<B>,
901        indices: IntTensor<B>,
902    ) -> MaxPool2dBackward<B>;
903
904    /// Down/up samples the input.
905    ///
906    /// # Shapes
907    ///
908    /// x: `[batch_size, channels, height, width]`,
909    fn interpolate(
910        x: FloatTensor<B>,
911        output_size: [usize; 2],
912        options: InterpolateOptions,
913    ) -> FloatTensor<B>;
914
915    /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
916    fn interpolate_backward(
917        x: FloatTensor<B>,
918        grad: FloatTensor<B>,
919        output_size: [usize; 2],
920        options: InterpolateOptions,
921    ) -> FloatTensor<B>;
922
923    /// Computes scaled dot-product attention: softmax(QKᵗ / √d) · V,
924    /// optionally applying a mask to the attention scores.
925    ///
926    /// # Arguments
927    /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q,  head_dim]`
928    /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
929    /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
930    /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
931    ///   here `true` indicates positions to mask (i.e. set to -∞ before softmax).
932    ///
933    /// # Returns
934    /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
935    /// representing the attended context per head.
936    ///
937    /// # Note
938    /// This implementation does not support dropout and is intended for inference or
939    /// use cases where dropout is not needed.
940    fn attention(
941        query: FloatTensor<B>,
942        key: FloatTensor<B>,
943        value: FloatTensor<B>,
944        mask: Option<BoolTensor<B>>,
945    ) -> FloatTensor<B> {
946        attention::naive_attention::<B>(query, key, value, mask)
947    }
948}
949
950#[cfg(test)]
951mod tests {
952    use super::*;
953
954    #[test]
955    #[should_panic = "stride must be non-zero"]
956    fn conv_options_stride_zero() {
957        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
958    }
959
960    #[test]
961    #[should_panic = "dilation must be non-zero"]
962    fn conv_options_dilation_zero() {
963        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
964    }
965
966    #[test]
967    #[should_panic = "groups must be non-zero"]
968    fn conv_options_groups_zero() {
969        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
970    }
971
972    #[test]
973    #[should_panic = "stride must be non-zero"]
974    fn conv_transpose_options_stride_zero() {
975        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
976    }
977
978    #[test]
979    #[should_panic = "dilation must be non-zero"]
980    fn conv_transpose_options_dilation_zero() {
981        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
982    }
983
984    #[test]
985    #[should_panic = "groups must be non-zero"]
986    fn conv_transpose_options_groups_zero() {
987        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
988    }
989
990    #[test]
991    #[should_panic = "stride must be non-zero"]
992    fn deform_conv_options_stride_zero() {
993        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
994    }
995
996    #[test]
997    #[should_panic = "dilation must be non-zero"]
998    fn deform_conv_options_dilation_zero() {
999        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
1000    }
1001
1002    #[test]
1003    #[should_panic = "weight groups must be non-zero"]
1004    fn deform_conv_options_weights_groups_zero() {
1005        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
1006    }
1007
1008    #[test]
1009    #[should_panic = "offset groups must be non-zero"]
1010    fn deform_conv_options_offset_groups_zero() {
1011        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
1012    }
1013
1014    #[test]
1015    #[should_panic = "stride must be non-zero"]
1016    fn unfold_options_stride_zero() {
1017        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
1018    }
1019
1020    #[test]
1021    #[should_panic = "dilation must be non-zero"]
1022    fn unfold_options_dilation_zero() {
1023        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
1024    }
1025}
burn_backend/backend/ops/modules/base.rs

burn_backend/backend/ops/modules/
base.rs
