burn_std/
ops.rs

1//! Configuration types for tensor operations.
2
3use crate::ElementConversion;
4use core::num::NonZeroUsize;
5
6/// Check that the parameter value is non-zero.
7// NOTE: for now we keep usize but we could refactor the parameters to hold `NonZeroUsize`.
8pub(crate) fn check_nonzero(value: usize, msg: &str) -> usize {
9    NonZeroUsize::new(value).expect(msg);
10    value
11}
12
13/// Convolution options.
14#[derive(Debug, Clone, Hash, PartialEq, Eq)]
15pub struct ConvOptions<const N: usize> {
16    /// Stride (non-zero).
17    pub stride: [usize; N],
18
19    /// Padding as `(begin, end)` pairs for each spatial dimension.
20    pub padding: [(usize, usize); N],
21
22    /// Dilation (non-zero).
23    pub dilation: [usize; N],
24
25    /// Groups (non-zero).
26    pub groups: usize,
27}
28
29impl<const N: usize> ConvOptions<N> {
30    /// Constructs a new `ConvOptions`.
31    pub fn new(
32        stride: [usize; N],
33        padding: [usize; N],
34        dilation: [usize; N],
35        groups: usize,
36    ) -> Self {
37        Self {
38            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
39            padding: padding.map(|padding| (padding, padding)),
40            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
41            groups: check_nonzero(groups, "groups must be non-zero"),
42        }
43    }
44
45    /// Constructs convolution options with explicit per-side padding.
46    pub fn new_with_padding(
47        stride: [usize; N],
48        padding: [(usize, usize); N],
49        dilation: [usize; N],
50        groups: usize,
51    ) -> Self {
52        Self {
53            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
54            padding,
55            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
56            groups: check_nonzero(groups, "groups must be non-zero"),
57        }
58    }
59
60    /// Returns true if padding is asymmetric.
61    pub fn is_asymmetric(&self) -> bool {
62        self.padding.iter().any(|(begin, end)| begin != end)
63    }
64
65    /// Returns the padding at the beginning of every spatial dimension.
66    pub fn padding_begin(&self) -> [usize; N] {
67        self.padding.map(|(begin, _)| begin)
68    }
69
70    /// Returns the padding at the end of every spatial dimension.
71    pub fn padding_end(&self) -> [usize; N] {
72        self.padding.map(|(_, after)| after)
73    }
74
75    /// Returns symmetric padding values.
76    ///
77    /// # Panics
78    /// Panics when any spatial dimension has asymmetric padding.
79    pub fn symmetric_padding(&self) -> [usize; N] {
80        assert!(
81            !self.is_asymmetric(),
82            "expected symmetric convolution padding"
83        );
84        self.padding_begin()
85    }
86}
87
88/// Convolution options with optional asymmetric end padding.
89///
90/// Deprecated compatibility wrapper for the convolution options API shipped in
91/// Burn 0.22. Use [`ConvOptions::new_with_padding`] instead.
92#[deprecated(since = "0.22.0", note = "Use `ConvOptions::new_with_padding` instead")]
93#[derive(Debug, Clone)]
94pub struct PaddedConvOptions<const N: usize> {
95    /// The underlying convolution options.
96    pub options: ConvOptions<N>,
97    /// Padding at the end of each spatial dimension.
98    ///
99    /// When `None`, the padding stored in [`Self::options`] is used unchanged.
100    pub padding_end: Option<[usize; N]>,
101}
102
103#[allow(deprecated)]
104impl<const N: usize> PaddedConvOptions<N> {
105    /// Creates options with explicit begin and end padding.
106    pub fn asymmetric(
107        stride: [usize; N],
108        padding_begin: [usize; N],
109        padding_end: [usize; N],
110        dilation: [usize; N],
111        groups: usize,
112    ) -> Self {
113        let options = ConvOptions::new(stride, padding_begin, dilation, groups);
114        let padding_end = (padding_begin != padding_end).then_some(padding_end);
115        Self {
116            options,
117            padding_end,
118        }
119    }
120
121    /// Returns true if explicit asymmetric end padding is present.
122    pub fn is_asymmetric(&self) -> bool {
123        self.padding_end.is_some()
124    }
125}
126
127#[allow(deprecated)]
128impl<const N: usize> From<PaddedConvOptions<N>> for ConvOptions<N> {
129    fn from(value: PaddedConvOptions<N>) -> Self {
130        let Some(padding_end) = value.padding_end else {
131            return value.options;
132        };
133        let padding_begin = value.options.padding_begin();
134        let padding = core::array::from_fn(|i| (padding_begin[i], padding_end[i]));
135
136        ConvOptions::new_with_padding(
137            value.options.stride,
138            padding,
139            value.options.dilation,
140            value.options.groups,
141        )
142    }
143}
144
145#[allow(deprecated)]
146impl<const N: usize> From<ConvOptions<N>> for PaddedConvOptions<N> {
147    fn from(options: ConvOptions<N>) -> Self {
148        if options.is_asymmetric() {
149            let padding_begin = options.padding_begin();
150            let padding_end = options.padding_end();
151            Self {
152                options: ConvOptions::new(
153                    options.stride,
154                    padding_begin,
155                    options.dilation,
156                    options.groups,
157                ),
158                padding_end: Some(padding_end),
159            }
160        } else {
161            Self {
162                options,
163                padding_end: None,
164            }
165        }
166    }
167}
168
169/// Max pooling options.
170#[derive(Debug, Clone, Hash, PartialEq, Eq)]
171pub struct MaxPoolOptions<const N: usize> {
172    /// Kernel size (non-zero).
173    pub kernel_size: [usize; N],
174
175    /// Stride (non-zero).
176    pub stride: [usize; N],
177
178    /// Padding as `(begin, end)` pairs for each spatial dimension.
179    pub padding: [(usize, usize); N],
180
181    /// Dilation (non-zero).
182    pub dilation: [usize; N],
183
184    /// If true, use ceiling instead of floor for output size calculation.
185    pub ceil_mode: bool,
186}
187
188impl<const N: usize> MaxPoolOptions<N> {
189    /// Constructs max pooling options with the given kernel size.
190    ///
191    /// Defaults to a stride equal to the kernel size, no padding, a dilation
192    /// of 1 and floor mode.
193    pub fn new(kernel_size: [usize; N]) -> Self {
194        let kernel_size = kernel_size.map(|k| check_nonzero(k, "kernel size must be non-zero"));
195        Self {
196            kernel_size,
197            stride: kernel_size,
198            padding: [(0, 0); N],
199            dilation: [1; N],
200            ceil_mode: false,
201        }
202    }
203
204    /// Sets the stride.
205    pub fn with_stride(mut self, stride: [usize; N]) -> Self {
206        self.stride = stride.map(|s| check_nonzero(s, "stride must be non-zero"));
207        self
208    }
209
210    /// Sets symmetric padding for every spatial dimension.
211    pub fn with_padding(mut self, padding: [usize; N]) -> Self {
212        self.padding = padding.map(|padding| (padding, padding));
213        self
214    }
215
216    /// Sets explicit `(begin, end)` padding for every spatial dimension.
217    pub fn with_padding_pairs(mut self, padding: [(usize, usize); N]) -> Self {
218        self.padding = padding;
219        self
220    }
221
222    /// Sets the dilation.
223    pub fn with_dilation(mut self, dilation: [usize; N]) -> Self {
224        self.dilation = dilation.map(|d| check_nonzero(d, "dilation must be non-zero"));
225        self
226    }
227
228    /// Sets whether to use ceiling instead of floor for output size calculation.
229    pub fn with_ceil_mode(mut self, ceil_mode: bool) -> Self {
230        self.ceil_mode = ceil_mode;
231        self
232    }
233
234    /// Returns true if padding is asymmetric.
235    pub fn is_asymmetric(&self) -> bool {
236        self.padding.iter().any(|(begin, end)| begin != end)
237    }
238}
239
240/// Average pooling options.
241#[derive(Debug, Clone, Hash, PartialEq, Eq)]
242pub struct AvgPoolOptions<const N: usize> {
243    /// Kernel size (non-zero).
244    pub kernel_size: [usize; N],
245
246    /// Stride (non-zero).
247    pub stride: [usize; N],
248
249    /// Padding as `(begin, end)` pairs for each spatial dimension.
250    pub padding: [(usize, usize); N],
251
252    /// Dilation (non-zero).
253    pub dilation: [usize; N],
254
255    /// If true, padded values are counted in the denominator of the average.
256    pub count_include_pad: bool,
257
258    /// If true, use ceiling instead of floor for output size calculation.
259    pub ceil_mode: bool,
260}
261
262impl<const N: usize> AvgPoolOptions<N> {
263    /// Constructs average pooling options with the given kernel size.
264    ///
265    /// Defaults to a stride equal to the kernel size, no padding, a dilation
266    /// of 1, counting padded values in the average and floor mode.
267    pub fn new(kernel_size: [usize; N]) -> Self {
268        let kernel_size = kernel_size.map(|k| check_nonzero(k, "kernel size must be non-zero"));
269        Self {
270            kernel_size,
271            stride: kernel_size,
272            padding: [(0, 0); N],
273            dilation: [1; N],
274            count_include_pad: true,
275            ceil_mode: false,
276        }
277    }
278
279    /// Sets the stride.
280    pub fn with_stride(mut self, stride: [usize; N]) -> Self {
281        self.stride = stride.map(|s| check_nonzero(s, "stride must be non-zero"));
282        self
283    }
284
285    /// Sets symmetric padding for every spatial dimension.
286    pub fn with_padding(mut self, padding: [usize; N]) -> Self {
287        self.padding = padding.map(|padding| (padding, padding));
288        self
289    }
290
291    /// Sets explicit `(begin, end)` padding for every spatial dimension.
292    pub fn with_padding_pairs(mut self, padding: [(usize, usize); N]) -> Self {
293        self.padding = padding;
294        self
295    }
296
297    /// Sets the dilation.
298    ///
299    /// Window element `k` reads input position `start + k * dilation`, as in max pooling.
300    pub fn with_dilation(mut self, dilation: [usize; N]) -> Self {
301        self.dilation = dilation.map(|d| check_nonzero(d, "dilation must be non-zero"));
302        self
303    }
304
305    /// Sets whether padded values are counted in the denominator of the average.
306    pub fn with_count_include_pad(mut self, count_include_pad: bool) -> Self {
307        self.count_include_pad = count_include_pad;
308        self
309    }
310
311    /// Sets whether to use ceiling instead of floor for output size calculation.
312    pub fn with_ceil_mode(mut self, ceil_mode: bool) -> Self {
313        self.ceil_mode = ceil_mode;
314        self
315    }
316
317    /// Returns true if padding is asymmetric.
318    pub fn is_asymmetric(&self) -> bool {
319        self.padding.iter().any(|(begin, end)| begin != end)
320    }
321}
322
323/// Deformable convolution options.
324#[derive(Debug, Clone, Hash, PartialEq, Eq)]
325pub struct DeformConvOptions<const N: usize> {
326    /// Stride (non-zero).
327    pub stride: [usize; N],
328
329    /// Padding.
330    pub padding: [usize; N],
331
332    /// Dilation (non-zero).
333    pub dilation: [usize; N],
334
335    /// Weight Groups (non-zero).
336    pub weight_groups: usize,
337
338    /// Offset Groups (non-zero).
339    pub offset_groups: usize,
340}
341
342impl<const N: usize> DeformConvOptions<N> {
343    /// Constructs a new `DeformConvOptions`.
344    pub fn new(
345        stride: [usize; N],
346        padding: [usize; N],
347        dilation: [usize; N],
348        weight_groups: usize,
349        offset_groups: usize,
350    ) -> Self {
351        Self {
352            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
353            padding,
354            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
355            weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"),
356            offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"),
357        }
358    }
359}
360
361/// Transposed convolution options.
362#[derive(Debug, Clone, Hash, PartialEq, Eq)]
363pub struct ConvTransposeOptions<const N: usize> {
364    /// Stride (non-zero).
365    pub stride: [usize; N],
366
367    /// Padding.
368    pub padding: [usize; N],
369
370    /// Padding out.
371    pub padding_out: [usize; N],
372
373    /// Dilation (non-zero).
374    pub dilation: [usize; N],
375
376    /// Groups (non-zero).
377    pub groups: usize,
378}
379
380impl<const N: usize> ConvTransposeOptions<N> {
381    /// Constructs a new `ConvTransposeOptions`.
382    pub fn new(
383        stride: [usize; N],
384        padding: [usize; N],
385        padding_out: [usize; N],
386        dilation: [usize; N],
387        groups: usize,
388    ) -> Self {
389        Self {
390            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
391            padding,
392            padding_out,
393            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
394            groups: check_nonzero(groups, "groups must be non-zero"),
395        }
396    }
397}
398
399/// Unfold operation options.
400#[derive(Debug, Clone)]
401pub struct UnfoldOptions {
402    /// The number of positions to slide over the input tensor in each dimension.
403    /// A stride of `[1, 1]` will slide the kernel one pixel at a time.
404    pub stride: [usize; 2],
405
406    /// The number of zero-padding pixels added to each side of the input tensor in each dimension.
407    pub padding: [usize; 2],
408
409    /// The spacing between the blocks (patches) in the original input tensor.
410    pub dilation: [usize; 2],
411}
412
413impl UnfoldOptions {
414    /// Constructs a new `UnfoldOptions`.
415    pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self {
416        Self {
417            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
418            padding,
419            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
420        }
421    }
422}
423
424/// Algorithm used.
425#[derive(new, Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
426pub enum InterpolateMode {
427    /// Nearest-neighbor floor interpolation.
428    /// Matches the legacy behavior of OpenCV’s INTER_NEAREST. It results in a bottom-right shift when resizing.
429    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
430    Nearest,
431
432    /// Nearest-neighbor exact interpolation.
433    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
434    NearestExact,
435
436    /// Bilinear interpolation.
437    /// <https://en.wikipedia.org/wiki/Bilinear_interpolation>
438    Bilinear,
439
440    /// Bicubic interpolation.
441    /// <https://en.wikipedia.org/wiki/Bicubic_interpolation>
442    Bicubic,
443
444    /// Lanczos3 interpolation (6-tap sinc-based filter).
445    /// <https://en.wikipedia.org/wiki/Lanczos_resampling>
446    Lanczos3,
447}
448
449/// Coordinate transformation mode for interpolation.
450#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
451pub enum CoordinateTransformMode {
452    /// Half-pixel coordinate mapping: `(x_out + 0.5) / scale - 0.5`.
453    HalfPixel,
454    /// Align corners coordinate mapping: `x_out * (l_in - 1) / (l_out - 1)`.
455    AlignCorners,
456    /// Asymmetric coordinate mapping: `x_out / scale`.
457    Asymmetric,
458    /// PyTorch half-pixel coordinate mapping: `(x_out + 0.5) / scale - 0.5` if `l_out > 1` else `0.0`.
459    PytorchHalfPixel,
460    /// Half-pixel symmetric coordinate mapping.
461    HalfPixelSymmetric,
462    /// TensorFlow crop and resize coordinate mapping with normalized ROI bounding box.
463    TfCropAndResize {
464        /// Normalized bounding box per spatial axis: [[y_start, y_end], [x_start, x_end]]
465        roi: [[f32; 2]; 2],
466        /// Extrapolation fill value for coordinates outside bounds
467        extrapolation_value: f32,
468    },
469}
470
471/// Nearest rounding mode for nearest neighbor interpolation.
472#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
473pub enum NearestRoundingMode {
474    /// Round with preference to floor at 0.5 boundaries (`ceil(x - 0.5)`).
475    RoundPreferFloor,
476    /// Round with preference to ceil at 0.5 boundaries (`floor(x + 0.5)`).
477    RoundPreferCeil,
478    /// Floor rounding (`floor(x)`).
479    Floor,
480    /// Ceil rounding (`ceil(x)`).
481    Ceil,
482}
483
484/// Configuration options for bicubic interpolation.
485#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
486pub struct CubicOptions {
487    /// Cubic spline parameter 'a'. Default: -0.75 (PyTorch/ONNX), -0.5 (TensorFlow).
488    pub a: f32,
489    /// Zero out-of-bounds tap weights and renormalize remaining weights.
490    pub exclude_outside: bool,
491}
492
493impl Default for CubicOptions {
494    fn default() -> Self {
495        Self {
496            a: -0.75,
497            exclude_outside: false,
498        }
499    }
500}
501
502/// Interpolation options.
503///
504/// Flex supports all coordinate transforms, rounding modes and cubic options.
505/// CubeCL supports the default nearest and nearest-exact combinations, plus bilinear,
506/// bicubic and Lanczos3 with `AlignCorners` or `HalfPixel` and default cubic options.
507/// With CubeCL, explicit scales must equal the resolved output/input ratios exactly,
508/// unless the transform is `AlignCorners`. Unsupported combinations panic.
509#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
510pub struct InterpolateOptions {
511    /// Algorithm used.
512    pub mode: InterpolateMode,
513    /// Coordinate transformation mode.
514    pub coordinate_transform: CoordinateTransformMode,
515    /// Nearest rounding mode.
516    pub nearest_mode: NearestRoundingMode,
517    /// Cubic interpolation options.
518    pub cubic_options: CubicOptions,
519    /// Output spatial size `[height, width]`.
520    ///
521    /// Exactly one of `output_size` or `scale_factor` must be set when calling
522    /// `burn_tensor::module::interpolate`.
523    pub output_size: Option<[usize; 2]>,
524    /// Scale factor `[scale_height, scale_width]` applied to the input spatial size.
525    /// The output size is `floor(input_size * scale_factor)`.
526    /// Scale-dependent transforms use these supplied factors for coordinates rather
527    /// than recomputing the output/input ratio. `AlignCorners` and `TfCropAndResize`
528    /// map coordinates from the resolved sizes instead.
529    pub scale_factor: Option<[f32; 2]>,
530}
531
532impl InterpolateOptions {
533    /// Create new interpolate options with mode-dependent defaults.
534    pub fn new(mode: InterpolateMode) -> Self {
535        let (coordinate_transform, nearest_mode) = match mode {
536            InterpolateMode::Nearest => (
537                CoordinateTransformMode::Asymmetric,
538                NearestRoundingMode::Floor,
539            ),
540            InterpolateMode::NearestExact => (
541                CoordinateTransformMode::HalfPixel,
542                NearestRoundingMode::RoundPreferCeil,
543            ),
544            InterpolateMode::Bilinear | InterpolateMode::Bicubic | InterpolateMode::Lanczos3 => (
545                CoordinateTransformMode::AlignCorners,
546                NearestRoundingMode::RoundPreferFloor,
547            ),
548        };
549
550        Self {
551            mode,
552            coordinate_transform,
553            nearest_mode,
554            cubic_options: CubicOptions::default(),
555            output_size: None,
556            scale_factor: None,
557        }
558    }
559
560    /// Set coordinate transformation mode.
561    pub fn with_coordinate_transform(mut self, mode: CoordinateTransformMode) -> Self {
562        self.coordinate_transform = mode;
563        self
564    }
565
566    /// Set nearest rounding mode.
567    pub fn with_nearest_mode(mut self, mode: NearestRoundingMode) -> Self {
568        self.nearest_mode = mode;
569        self
570    }
571
572    /// Set cubic options.
573    pub fn with_cubic_options(mut self, options: CubicOptions) -> Self {
574        self.cubic_options = options;
575        self
576    }
577
578    /// Set corner alignment for bilinear, bicubic and Lanczos3 interpolation.
579    /// `true` selects `AlignCorners`; `false` selects `HalfPixel`.
580    /// Nearest modes are unchanged; use [`Self::with_coordinate_transform`] to set their mapping.
581    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
582        if !matches!(
583            self.mode,
584            InterpolateMode::Nearest | InterpolateMode::NearestExact
585        ) {
586            self.coordinate_transform = if align_corners {
587                CoordinateTransformMode::AlignCorners
588            } else {
589                CoordinateTransformMode::HalfPixel
590            };
591        }
592        self
593    }
594
595    /// Returns `true` if coordinate transform is `AlignCorners`.
596    pub fn align_corners(&self) -> bool {
597        matches!(
598            self.coordinate_transform,
599            CoordinateTransformMode::AlignCorners
600        )
601    }
602
603    /// Set the output size `[height, width]`, clearing any scale factor.
604    pub fn with_output_size(mut self, output_size: [usize; 2]) -> Self {
605        self.output_size = Some(output_size);
606        self.scale_factor = None;
607        self
608    }
609
610    /// Set the scale factor `[scale_height, scale_width]`, clearing any output size.
611    pub fn with_scale_factor(mut self, scale_factor: [f32; 2]) -> Self {
612        self.scale_factor = Some(scale_factor);
613        self.output_size = None;
614        self
615    }
616}
617
618/// Padding mode for grid sampling when coordinates are out of bounds.
619///
620/// Matches PyTorch's `padding_mode` parameter in `grid_sample`.
621#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)]
622pub enum GridSamplePaddingMode {
623    /// Fill with zeros for out-of-bounds coordinates.
624    #[default]
625    Zeros,
626    /// Clamp coordinates to the border (use nearest edge value).
627    Border,
628    /// Reflect coordinates at the boundary.
629    Reflection,
630}
631
632/// Options for grid sampling operations.
633#[derive(Debug, Clone)]
634pub struct GridSampleOptions {
635    /// Interpolation mode (bilinear, nearest, or bicubic).
636    pub mode: InterpolateMode,
637    /// Padding mode for out-of-bounds coordinates.
638    pub padding_mode: GridSamplePaddingMode,
639    /// If `true`, grid values of -1 and 1 correspond to the corner pixels.
640    /// If `false`, they correspond to the corner points of the corner pixels
641    /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates).
642    pub align_corners: bool,
643}
644
645impl Default for GridSampleOptions {
646    fn default() -> Self {
647        Self {
648            mode: InterpolateMode::Bilinear,
649            padding_mode: GridSamplePaddingMode::Zeros,
650            align_corners: false,
651        }
652    }
653}
654
655impl From<InterpolateMode> for GridSampleOptions {
656    fn from(value: InterpolateMode) -> Self {
657        GridSampleOptions::new(value)
658    }
659}
660
661impl GridSampleOptions {
662    /// Create new grid sample options with the given interpolation mode.
663    ///
664    /// Uses default values for padding_mode (Zeros) and align_corners (false).
665    pub fn new(mode: InterpolateMode) -> Self {
666        Self {
667            mode,
668            ..Default::default()
669        }
670    }
671
672    /// Set the padding mode.
673    pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self {
674        self.padding_mode = padding_mode;
675        self
676    }
677
678    /// Set align_corners.
679    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
680        self.align_corners = align_corners;
681        self
682    }
683}
684
685/// Padding mode for tensor pad operations.
686///
687/// Defines how values are filled when padding a tensor beyond its original boundaries.
688/// Padding can be applied to any dimension of a tensor.
689///
690/// # Modes
691///
692/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0)
693/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size)
694/// - [`Edge`](PadMode::Edge): Replicate boundary values
695#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)]
696pub enum PadMode {
697    /// Fill padded regions with a constant value.
698    ///
699    /// # Example
700    /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0:
701    /// Result: `[0, 0, 1, 2, 3]`
702    Constant(f32),
703
704    /// Reflect values at the boundary, excluding the edge value.
705    ///
706    /// Padding must be less than the dimension size (i.e., `padding < dim_size`).
707    ///
708    /// # Example
709    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
710    /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0)
711    Reflect,
712
713    /// Replicate the edge values.
714    ///
715    /// # Example
716    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
717    /// Result: `[1, 1, 1, 2, 3, 4]`
718    Edge,
719}
720
721impl Default for PadMode {
722    fn default() -> Self {
723        PadMode::Constant(0.0)
724    }
725}
726
727impl<E: ElementConversion> From<E> for PadMode {
728    fn from(value: E) -> Self {
729        PadMode::Constant(value.elem())
730    }
731}
732
733/// Options for the attention module.
734#[derive(Debug, Clone, Copy, Default, PartialEq, serde::Deserialize, serde::Serialize)]
735pub struct AttentionModuleOptions {
736    /// Custom scale factor applied to QK^T. When `None`, defaults to `1/sqrt(head_dim)`.
737    pub scale: Option<f64>,
738
739    /// Soft capping applied before softmax: `softcap * tanh(scores / softcap)`.
740    /// Used by Gemma-2 and similar models. Must be positive when set.
741    pub softcap: Option<f64>,
742
743    /// When `true`, applies causal (autoregressive) masking so that each query position
744    /// can only attend to key positions at or before it. This is more efficient than
745    /// passing an explicit lower-triangular bool mask because backends can use optimized
746    /// kernel paths (e.g. flash attention with causal mode).
747    ///
748    /// When `seq_q != seq_k`, [`causal_alignment`](Self::causal_alignment) decides which
749    /// query row lines up with which key column.
750    pub is_causal: bool,
751
752    /// Where the causal diagonal is anchored when `seq_q != seq_k`. Ignored unless
753    /// `is_causal` is set. Both alignments agree when `seq_q == seq_k`.
754    ///
755    /// Options serialized before this field existed deserialize as
756    /// [`BottomRight`](CausalAlignment::BottomRight), the alignment burn used until then.
757    #[serde(default)]
758    pub causal_alignment: CausalAlignment,
759}
760
761/// Anchor of the causal diagonal for [`AttentionModuleOptions::is_causal`].
762///
763/// Query `i` may attend key `j` iff `j <= i + offset`, where `offset` is `0` for
764/// [`TopLeft`](Self::TopLeft) and `seq_k - seq_q` for [`BottomRight`](Self::BottomRight).
765#[derive(
766    Debug, Clone, Copy, Default, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
767)]
768pub enum CausalAlignment {
769    /// The last query row sees every key: `offset = seq_k - seq_q`. This is the
770    /// KV-cache convention (the queries are the newest `seq_q` positions), used by
771    /// FlashAttention 2.1+ and by ONNX `Attention` when a `past_key` is supplied.
772    #[default]
773    BottomRight,
774    /// The first query row sees the first key: `offset = 0` (`tril` on `[seq_q, seq_k]`).
775    /// This is PyTorch's `scaled_dot_product_attention(is_causal=True)` and ONNX
776    /// `Attention` without a KV cache.
777    TopLeft,
778}
779
780impl CausalAlignment {
781    /// The offset of the causal diagonal: query `i` may attend key `j` iff `j <= i + offset`.
782    pub fn offset(self, seq_q: usize, seq_k: usize) -> i64 {
783        match self {
784            CausalAlignment::BottomRight => seq_k as i64 - seq_q as i64,
785            CausalAlignment::TopLeft => 0,
786        }
787    }
788}
789
790/// Computation to be used to update the existing values in indexed assignment operations (scatter/select).
791#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
792pub enum IndexingUpdateOp {
793    /// Overwrite existing values.
794    Assign,
795    /// Performs an addition.
796    Add,
797    /// Multiply existing values.
798    Mul,
799    /// Take element-wise minimum.
800    Min,
801    /// Take element-wise maximum.
802    Max,
803}
804
805#[cfg(test)]
806mod tests {
807    use super::*;
808
809    #[test]
810    fn conv_options_symmetric_constructor() {
811        let options = ConvOptions::new([1, 2], [3, 4], [5, 6], 7);
812
813        assert_eq!(options.padding, [(3, 3), (4, 4)]);
814        assert_eq!(options.padding_begin(), [3, 4]);
815        assert_eq!(options.padding_end(), [3, 4]);
816        assert_eq!(options.symmetric_padding(), [3, 4]);
817        assert!(!options.is_asymmetric());
818    }
819
820    #[test]
821    fn conv_options_explicit_padding_constructor() {
822        let options = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
823
824        assert_eq!(options.padding, [(3, 4), (5, 6)]);
825        assert_eq!(options.padding_begin(), [3, 5]);
826        assert_eq!(options.padding_end(), [4, 6]);
827        assert!(options.is_asymmetric());
828    }
829
830    #[test]
831    #[allow(deprecated)]
832    fn padded_conv_options_convert_to_conv_options() {
833        let options: ConvOptions<2> =
834            PaddedConvOptions::asymmetric([1, 2], [3, 5], [4, 6], [7, 8], 9).into();
835
836        assert_eq!(options.padding, [(3, 4), (5, 6)]);
837        assert_eq!(options.stride, [1, 2]);
838        assert_eq!(options.dilation, [7, 8]);
839        assert_eq!(options.groups, 9);
840    }
841
842    #[test]
843    #[allow(deprecated)]
844    fn asymmetric_conv_options_roundtrip_through_padded_options() {
845        let expected = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
846        let padded: PaddedConvOptions<2> = expected.clone().into();
847        let actual: ConvOptions<2> = padded.into();
848
849        assert_eq!(actual, expected);
850    }
851
852    #[test]
853    #[should_panic = "expected symmetric convolution padding"]
854    fn conv_options_symmetric_padding_with_asymmetric_options() {
855        let options = ConvOptions::new_with_padding([1], [(1, 2)], [1], 1);
856        let _ = options.symmetric_padding();
857    }
858
859    #[test]
860    #[should_panic = "stride must be non-zero"]
861    fn conv_options_stride_zero() {
862        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
863    }
864
865    #[test]
866    #[should_panic = "dilation must be non-zero"]
867    fn conv_options_dilation_zero() {
868        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
869    }
870
871    #[test]
872    #[should_panic = "groups must be non-zero"]
873    fn conv_options_groups_zero() {
874        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
875    }
876
877    #[test]
878    fn pool_options_defaults() {
879        let max = MaxPoolOptions::new([2, 3]);
880        let avg = AvgPoolOptions::new([2, 3]);
881
882        assert_eq!(max.stride, [2, 3]);
883        assert_eq!(max.padding, [(0, 0), (0, 0)]);
884        assert_eq!(max.dilation, [1, 1]);
885        assert!(!max.ceil_mode);
886        assert_eq!(avg.stride, [2, 3]);
887        assert_eq!(avg.padding, [(0, 0), (0, 0)]);
888        assert_eq!(avg.dilation, [1, 1]);
889        assert!(avg.count_include_pad);
890        assert!(!avg.ceil_mode);
891    }
892
893    #[test]
894    fn pool_options_padding() {
895        let symmetric = MaxPoolOptions::new([2, 2]).with_padding([0, 1]);
896        let asymmetric = AvgPoolOptions::new([2]).with_padding_pairs([(0, 1)]);
897
898        assert_eq!(symmetric.padding, [(0, 0), (1, 1)]);
899        assert!(!symmetric.is_asymmetric());
900        assert!(asymmetric.is_asymmetric());
901    }
902
903    #[test]
904    #[should_panic = "kernel size must be non-zero"]
905    fn max_pool_options_kernel_size_zero() {
906        let _opt = MaxPoolOptions::new([0, 1]);
907    }
908
909    #[test]
910    #[should_panic = "dilation must be non-zero"]
911    fn max_pool_options_dilation_zero() {
912        let _opt = MaxPoolOptions::new([1, 1]).with_dilation([0, 1]);
913    }
914
915    #[test]
916    #[should_panic = "dilation must be non-zero"]
917    fn avg_pool_options_dilation_zero() {
918        let _opt = AvgPoolOptions::new([1, 1]).with_dilation([0, 1]);
919    }
920
921    #[test]
922    #[should_panic = "stride must be non-zero"]
923    fn avg_pool_options_stride_zero() {
924        let _opt = AvgPoolOptions::new([1, 1]).with_stride([0, 1]);
925    }
926
927    #[test]
928    #[should_panic = "stride must be non-zero"]
929    fn conv_transpose_options_stride_zero() {
930        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
931    }
932
933    #[test]
934    #[should_panic = "dilation must be non-zero"]
935    fn conv_transpose_options_dilation_zero() {
936        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
937    }
938
939    #[test]
940    #[should_panic = "groups must be non-zero"]
941    fn conv_transpose_options_groups_zero() {
942        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
943    }
944
945    #[test]
946    #[should_panic = "stride must be non-zero"]
947    fn deform_conv_options_stride_zero() {
948        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
949    }
950
951    #[test]
952    #[should_panic = "dilation must be non-zero"]
953    fn deform_conv_options_dilation_zero() {
954        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
955    }
956
957    #[test]
958    #[should_panic = "weight groups must be non-zero"]
959    fn deform_conv_options_weights_groups_zero() {
960        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
961    }
962
963    #[test]
964    #[should_panic = "offset groups must be non-zero"]
965    fn deform_conv_options_offset_groups_zero() {
966        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
967    }
968
969    #[test]
970    #[should_panic = "stride must be non-zero"]
971    fn unfold_options_stride_zero() {
972        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
973    }
974
975    #[test]
976    #[should_panic = "dilation must be non-zero"]
977    fn unfold_options_dilation_zero() {
978        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
979    }
980}
burn_std/ops.rs

burn_std/
ops.rs
