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/// Deformable convolution options.
170#[derive(Debug, Clone, Hash, PartialEq, Eq)]
171pub struct DeformConvOptions<const N: usize> {
172    /// Stride (non-zero).
173    pub stride: [usize; N],
174
175    /// Padding.
176    pub padding: [usize; N],
177
178    /// Dilation (non-zero).
179    pub dilation: [usize; N],
180
181    /// Weight Groups (non-zero).
182    pub weight_groups: usize,
183
184    /// Offset Groups (non-zero).
185    pub offset_groups: usize,
186}
187
188impl<const N: usize> DeformConvOptions<N> {
189    /// Constructs a new `DeformConvOptions`.
190    pub fn new(
191        stride: [usize; N],
192        padding: [usize; N],
193        dilation: [usize; N],
194        weight_groups: usize,
195        offset_groups: usize,
196    ) -> Self {
197        Self {
198            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
199            padding,
200            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
201            weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"),
202            offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"),
203        }
204    }
205}
206
207/// Transposed convolution options.
208#[derive(Debug, Clone, Hash, PartialEq, Eq)]
209pub struct ConvTransposeOptions<const N: usize> {
210    /// Stride (non-zero).
211    pub stride: [usize; N],
212
213    /// Padding.
214    pub padding: [usize; N],
215
216    /// Padding out.
217    pub padding_out: [usize; N],
218
219    /// Dilation (non-zero).
220    pub dilation: [usize; N],
221
222    /// Groups (non-zero).
223    pub groups: usize,
224}
225
226impl<const N: usize> ConvTransposeOptions<N> {
227    /// Constructs a new `ConvTransposeOptions`.
228    pub fn new(
229        stride: [usize; N],
230        padding: [usize; N],
231        padding_out: [usize; N],
232        dilation: [usize; N],
233        groups: usize,
234    ) -> Self {
235        Self {
236            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
237            padding,
238            padding_out,
239            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
240            groups: check_nonzero(groups, "groups must be non-zero"),
241        }
242    }
243}
244
245/// Unfold operation options.
246#[derive(Debug, Clone)]
247pub struct UnfoldOptions {
248    /// The number of positions to slide over the input tensor in each dimension.
249    /// A stride of `[1, 1]` will slide the kernel one pixel at a time.
250    pub stride: [usize; 2],
251
252    /// The number of zero-padding pixels added to each side of the input tensor in each dimension.
253    pub padding: [usize; 2],
254
255    /// The spacing between the blocks (patches) in the original input tensor.
256    pub dilation: [usize; 2],
257}
258
259impl UnfoldOptions {
260    /// Constructs a new `UnfoldOptions`.
261    pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self {
262        Self {
263            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
264            padding,
265            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
266        }
267    }
268}
269
270/// Algorithm used.
271#[derive(new, Debug, Clone, serde::Deserialize, serde::Serialize)]
272pub enum InterpolateMode {
273    /// Nearest-neighbor floor interpolation.
274    /// Matches the legacy behavior of OpenCV’s INTER_NEAREST. It results in a bottom-right shift when resizing.
275    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
276    Nearest,
277
278    /// Nearest-neighbor exact interpolation.
279    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
280    NearestExact,
281
282    /// Bilinear interpolation.
283    /// <https://en.wikipedia.org/wiki/Bilinear_interpolation>
284    Bilinear,
285
286    /// Bicubic interpolation.
287    /// <https://en.wikipedia.org/wiki/Bicubic_interpolation>
288    Bicubic,
289
290    /// Lanczos3 interpolation (6-tap sinc-based filter).
291    /// <https://en.wikipedia.org/wiki/Lanczos_resampling>
292    Lanczos3,
293}
294
295/// Interpolation options.
296#[derive(Debug, Clone)]
297pub struct InterpolateOptions {
298    /// Algorithm used.
299    pub mode: InterpolateMode,
300    /// If `true`, the input and output tensors are aligned by their corner pixels.
301    /// If `false`, half-pixel coordinate mapping is used instead.
302    pub align_corners: bool,
303}
304
305impl InterpolateOptions {
306    /// Create new interpolate options with the given mode.
307    /// Defaults to `align_corners = true`.
308    pub fn new(mode: InterpolateMode) -> Self {
309        Self {
310            mode,
311            align_corners: true,
312        }
313    }
314
315    /// Set align_corners.
316    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
317        self.align_corners = align_corners;
318        self
319    }
320}
321
322/// Padding mode for grid sampling when coordinates are out of bounds.
323///
324/// Matches PyTorch's `padding_mode` parameter in `grid_sample`.
325#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)]
326pub enum GridSamplePaddingMode {
327    /// Fill with zeros for out-of-bounds coordinates.
328    #[default]
329    Zeros,
330    /// Clamp coordinates to the border (use nearest edge value).
331    Border,
332    /// Reflect coordinates at the boundary.
333    Reflection,
334}
335
336/// Options for grid sampling operations.
337#[derive(Debug, Clone)]
338pub struct GridSampleOptions {
339    /// Interpolation mode (bilinear, nearest, or bicubic).
340    pub mode: InterpolateMode,
341    /// Padding mode for out-of-bounds coordinates.
342    pub padding_mode: GridSamplePaddingMode,
343    /// If `true`, grid values of -1 and 1 correspond to the corner pixels.
344    /// If `false`, they correspond to the corner points of the corner pixels
345    /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates).
346    pub align_corners: bool,
347}
348
349impl Default for GridSampleOptions {
350    fn default() -> Self {
351        Self {
352            mode: InterpolateMode::Bilinear,
353            padding_mode: GridSamplePaddingMode::Zeros,
354            align_corners: false,
355        }
356    }
357}
358
359impl From<InterpolateMode> for GridSampleOptions {
360    fn from(value: InterpolateMode) -> Self {
361        GridSampleOptions::new(value)
362    }
363}
364
365impl GridSampleOptions {
366    /// Create new grid sample options with the given interpolation mode.
367    ///
368    /// Uses default values for padding_mode (Zeros) and align_corners (false).
369    pub fn new(mode: InterpolateMode) -> Self {
370        Self {
371            mode,
372            ..Default::default()
373        }
374    }
375
376    /// Set the padding mode.
377    pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self {
378        self.padding_mode = padding_mode;
379        self
380    }
381
382    /// Set align_corners.
383    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
384        self.align_corners = align_corners;
385        self
386    }
387}
388
389/// Padding mode for tensor pad operations.
390///
391/// Defines how values are filled when padding a tensor beyond its original boundaries.
392/// Padding can be applied to any dimension of a tensor.
393///
394/// # Modes
395///
396/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0)
397/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size)
398/// - [`Edge`](PadMode::Edge): Replicate boundary values
399#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)]
400pub enum PadMode {
401    /// Fill padded regions with a constant value.
402    ///
403    /// # Example
404    /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0:
405    /// Result: `[0, 0, 1, 2, 3]`
406    Constant(f32),
407
408    /// Reflect values at the boundary, excluding the edge value.
409    ///
410    /// Padding must be less than the dimension size (i.e., `padding < dim_size`).
411    ///
412    /// # Example
413    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
414    /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0)
415    Reflect,
416
417    /// Replicate the edge values.
418    ///
419    /// # Example
420    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
421    /// Result: `[1, 1, 1, 2, 3, 4]`
422    Edge,
423}
424
425impl Default for PadMode {
426    fn default() -> Self {
427        PadMode::Constant(0.0)
428    }
429}
430
431impl<E: ElementConversion> From<E> for PadMode {
432    fn from(value: E) -> Self {
433        PadMode::Constant(value.elem())
434    }
435}
436
437/// Options for the attention module.
438#[derive(Debug, Clone, Copy, Default, PartialEq, serde::Deserialize, serde::Serialize)]
439pub struct AttentionModuleOptions {
440    /// Custom scale factor applied to QK^T. When `None`, defaults to `1/sqrt(head_dim)`.
441    pub scale: Option<f64>,
442
443    /// Soft capping applied before softmax: `softcap * tanh(scores / softcap)`.
444    /// Used by Gemma-2 and similar models. Must be positive when set.
445    pub softcap: Option<f64>,
446
447    /// When `true`, applies causal (autoregressive) masking so that each query position
448    /// can only attend to key positions at or before it. This is more efficient than
449    /// passing an explicit lower-triangular bool mask because backends can use optimized
450    /// kernel paths (e.g. flash attention with causal mode).
451    pub is_causal: bool,
452}
453
454/// Computation to be used to update the existing values in indexed assignment operations (scatter/select).
455#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
456pub enum IndexingUpdateOp {
457    /// Overwrite existing values.
458    Assign,
459    /// Performs an addition.
460    Add,
461    /// Multiply existing values.
462    Mul,
463    /// Take element-wise minimum.
464    Min,
465    /// Take element-wise maximum.
466    Max,
467}
468
469#[cfg(test)]
470mod tests {
471    use super::*;
472
473    #[test]
474    fn conv_options_symmetric_constructor() {
475        let options = ConvOptions::new([1, 2], [3, 4], [5, 6], 7);
476
477        assert_eq!(options.padding, [(3, 3), (4, 4)]);
478        assert_eq!(options.padding_begin(), [3, 4]);
479        assert_eq!(options.padding_end(), [3, 4]);
480        assert_eq!(options.symmetric_padding(), [3, 4]);
481        assert!(!options.is_asymmetric());
482    }
483
484    #[test]
485    fn conv_options_explicit_padding_constructor() {
486        let options = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
487
488        assert_eq!(options.padding, [(3, 4), (5, 6)]);
489        assert_eq!(options.padding_begin(), [3, 5]);
490        assert_eq!(options.padding_end(), [4, 6]);
491        assert!(options.is_asymmetric());
492    }
493
494    #[test]
495    #[allow(deprecated)]
496    fn padded_conv_options_convert_to_conv_options() {
497        let options: ConvOptions<2> =
498            PaddedConvOptions::asymmetric([1, 2], [3, 5], [4, 6], [7, 8], 9).into();
499
500        assert_eq!(options.padding, [(3, 4), (5, 6)]);
501        assert_eq!(options.stride, [1, 2]);
502        assert_eq!(options.dilation, [7, 8]);
503        assert_eq!(options.groups, 9);
504    }
505
506    #[test]
507    #[allow(deprecated)]
508    fn asymmetric_conv_options_roundtrip_through_padded_options() {
509        let expected = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
510        let padded: PaddedConvOptions<2> = expected.clone().into();
511        let actual: ConvOptions<2> = padded.into();
512
513        assert_eq!(actual, expected);
514    }
515
516    #[test]
517    #[should_panic = "expected symmetric convolution padding"]
518    fn conv_options_symmetric_padding_with_asymmetric_options() {
519        let options = ConvOptions::new_with_padding([1], [(1, 2)], [1], 1);
520        let _ = options.symmetric_padding();
521    }
522
523    #[test]
524    #[should_panic = "stride must be non-zero"]
525    fn conv_options_stride_zero() {
526        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
527    }
528
529    #[test]
530    #[should_panic = "dilation must be non-zero"]
531    fn conv_options_dilation_zero() {
532        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
533    }
534
535    #[test]
536    #[should_panic = "groups must be non-zero"]
537    fn conv_options_groups_zero() {
538        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
539    }
540
541    #[test]
542    #[should_panic = "stride must be non-zero"]
543    fn conv_transpose_options_stride_zero() {
544        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
545    }
546
547    #[test]
548    #[should_panic = "dilation must be non-zero"]
549    fn conv_transpose_options_dilation_zero() {
550        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
551    }
552
553    #[test]
554    #[should_panic = "groups must be non-zero"]
555    fn conv_transpose_options_groups_zero() {
556        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
557    }
558
559    #[test]
560    #[should_panic = "stride must be non-zero"]
561    fn deform_conv_options_stride_zero() {
562        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
563    }
564
565    #[test]
566    #[should_panic = "dilation must be non-zero"]
567    fn deform_conv_options_dilation_zero() {
568        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
569    }
570
571    #[test]
572    #[should_panic = "weight groups must be non-zero"]
573    fn deform_conv_options_weights_groups_zero() {
574        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
575    }
576
577    #[test]
578    #[should_panic = "offset groups must be non-zero"]
579    fn deform_conv_options_offset_groups_zero() {
580        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
581    }
582
583    #[test]
584    #[should_panic = "stride must be non-zero"]
585    fn unfold_options_stride_zero() {
586        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
587    }
588
589    #[test]
590    #[should_panic = "dilation must be non-zero"]
591    fn unfold_options_dilation_zero() {
592        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
593    }
594}
burn_std/ops.rs

burn_std/
ops.rs
