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    /// If true, padded values are counted in the denominator of the average.
253    pub count_include_pad: bool,
254
255    /// If true, use ceiling instead of floor for output size calculation.
256    pub ceil_mode: bool,
257}
258
259impl<const N: usize> AvgPoolOptions<N> {
260    /// Constructs average pooling options with the given kernel size.
261    ///
262    /// Defaults to a stride equal to the kernel size, no padding, counting
263    /// padded values in the average and floor mode.
264    pub fn new(kernel_size: [usize; N]) -> Self {
265        let kernel_size = kernel_size.map(|k| check_nonzero(k, "kernel size must be non-zero"));
266        Self {
267            kernel_size,
268            stride: kernel_size,
269            padding: [(0, 0); N],
270            count_include_pad: true,
271            ceil_mode: false,
272        }
273    }
274
275    /// Sets the stride.
276    pub fn with_stride(mut self, stride: [usize; N]) -> Self {
277        self.stride = stride.map(|s| check_nonzero(s, "stride must be non-zero"));
278        self
279    }
280
281    /// Sets symmetric padding for every spatial dimension.
282    pub fn with_padding(mut self, padding: [usize; N]) -> Self {
283        self.padding = padding.map(|padding| (padding, padding));
284        self
285    }
286
287    /// Sets explicit `(begin, end)` padding for every spatial dimension.
288    pub fn with_padding_pairs(mut self, padding: [(usize, usize); N]) -> Self {
289        self.padding = padding;
290        self
291    }
292
293    /// Sets whether padded values are counted in the denominator of the average.
294    pub fn with_count_include_pad(mut self, count_include_pad: bool) -> Self {
295        self.count_include_pad = count_include_pad;
296        self
297    }
298
299    /// Sets whether to use ceiling instead of floor for output size calculation.
300    pub fn with_ceil_mode(mut self, ceil_mode: bool) -> Self {
301        self.ceil_mode = ceil_mode;
302        self
303    }
304
305    /// Returns true if padding is asymmetric.
306    pub fn is_asymmetric(&self) -> bool {
307        self.padding.iter().any(|(begin, end)| begin != end)
308    }
309}
310
311/// Deformable convolution options.
312#[derive(Debug, Clone, Hash, PartialEq, Eq)]
313pub struct DeformConvOptions<const N: usize> {
314    /// Stride (non-zero).
315    pub stride: [usize; N],
316
317    /// Padding.
318    pub padding: [usize; N],
319
320    /// Dilation (non-zero).
321    pub dilation: [usize; N],
322
323    /// Weight Groups (non-zero).
324    pub weight_groups: usize,
325
326    /// Offset Groups (non-zero).
327    pub offset_groups: usize,
328}
329
330impl<const N: usize> DeformConvOptions<N> {
331    /// Constructs a new `DeformConvOptions`.
332    pub fn new(
333        stride: [usize; N],
334        padding: [usize; N],
335        dilation: [usize; N],
336        weight_groups: usize,
337        offset_groups: usize,
338    ) -> Self {
339        Self {
340            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
341            padding,
342            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
343            weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"),
344            offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"),
345        }
346    }
347}
348
349/// Transposed convolution options.
350#[derive(Debug, Clone, Hash, PartialEq, Eq)]
351pub struct ConvTransposeOptions<const N: usize> {
352    /// Stride (non-zero).
353    pub stride: [usize; N],
354
355    /// Padding.
356    pub padding: [usize; N],
357
358    /// Padding out.
359    pub padding_out: [usize; N],
360
361    /// Dilation (non-zero).
362    pub dilation: [usize; N],
363
364    /// Groups (non-zero).
365    pub groups: usize,
366}
367
368impl<const N: usize> ConvTransposeOptions<N> {
369    /// Constructs a new `ConvTransposeOptions`.
370    pub fn new(
371        stride: [usize; N],
372        padding: [usize; N],
373        padding_out: [usize; N],
374        dilation: [usize; N],
375        groups: usize,
376    ) -> Self {
377        Self {
378            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
379            padding,
380            padding_out,
381            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
382            groups: check_nonzero(groups, "groups must be non-zero"),
383        }
384    }
385}
386
387/// Unfold operation options.
388#[derive(Debug, Clone)]
389pub struct UnfoldOptions {
390    /// The number of positions to slide over the input tensor in each dimension.
391    /// A stride of `[1, 1]` will slide the kernel one pixel at a time.
392    pub stride: [usize; 2],
393
394    /// The number of zero-padding pixels added to each side of the input tensor in each dimension.
395    pub padding: [usize; 2],
396
397    /// The spacing between the blocks (patches) in the original input tensor.
398    pub dilation: [usize; 2],
399}
400
401impl UnfoldOptions {
402    /// Constructs a new `UnfoldOptions`.
403    pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self {
404        Self {
405            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
406            padding,
407            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
408        }
409    }
410}
411
412/// Algorithm used.
413#[derive(new, Debug, Clone, serde::Deserialize, serde::Serialize)]
414pub enum InterpolateMode {
415    /// Nearest-neighbor floor interpolation.
416    /// Matches the legacy behavior of OpenCV’s INTER_NEAREST. It results in a bottom-right shift when resizing.
417    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
418    Nearest,
419
420    /// Nearest-neighbor exact interpolation.
421    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
422    NearestExact,
423
424    /// Bilinear interpolation.
425    /// <https://en.wikipedia.org/wiki/Bilinear_interpolation>
426    Bilinear,
427
428    /// Bicubic interpolation.
429    /// <https://en.wikipedia.org/wiki/Bicubic_interpolation>
430    Bicubic,
431
432    /// Lanczos3 interpolation (6-tap sinc-based filter).
433    /// <https://en.wikipedia.org/wiki/Lanczos_resampling>
434    Lanczos3,
435}
436
437/// Interpolation options.
438#[derive(Debug, Clone)]
439pub struct InterpolateOptions {
440    /// Algorithm used.
441    pub mode: InterpolateMode,
442    /// If `true`, the input and output tensors are aligned by their corner pixels.
443    /// If `false`, half-pixel coordinate mapping is used instead.
444    pub align_corners: bool,
445    /// Output spatial size `[height, width]`.
446    ///
447    /// Exactly one of `output_size` or `scale_factor` must be set when calling
448    /// `burn_tensor::module::interpolate`. Backends receive the resolved size as a
449    /// separate argument and ignore this field.
450    pub output_size: Option<[usize; 2]>,
451    /// Scale factor `[scale_height, scale_width]` applied to the input spatial size.
452    /// The output size is `floor(input_size * scale_factor)`.
453    ///
454    /// Exactly one of `output_size` or `scale_factor` must be set when calling
455    /// `burn_tensor::module::interpolate`. Backends ignore this field.
456    pub scale_factor: Option<[f32; 2]>,
457}
458
459impl InterpolateOptions {
460    /// Create new interpolate options with the given mode.
461    /// Defaults to `align_corners = true` and no output size or scale factor.
462    pub fn new(mode: InterpolateMode) -> Self {
463        Self {
464            mode,
465            align_corners: true,
466            output_size: None,
467            scale_factor: None,
468        }
469    }
470
471    /// Set align_corners.
472    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
473        self.align_corners = align_corners;
474        self
475    }
476
477    /// Set the output size `[height, width]`, clearing any scale factor.
478    pub fn with_output_size(mut self, output_size: [usize; 2]) -> Self {
479        self.output_size = Some(output_size);
480        self.scale_factor = None;
481        self
482    }
483
484    /// Set the scale factor `[scale_height, scale_width]`, clearing any output size.
485    pub fn with_scale_factor(mut self, scale_factor: [f32; 2]) -> Self {
486        self.scale_factor = Some(scale_factor);
487        self.output_size = None;
488        self
489    }
490}
491
492/// Padding mode for grid sampling when coordinates are out of bounds.
493///
494/// Matches PyTorch's `padding_mode` parameter in `grid_sample`.
495#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)]
496pub enum GridSamplePaddingMode {
497    /// Fill with zeros for out-of-bounds coordinates.
498    #[default]
499    Zeros,
500    /// Clamp coordinates to the border (use nearest edge value).
501    Border,
502    /// Reflect coordinates at the boundary.
503    Reflection,
504}
505
506/// Options for grid sampling operations.
507#[derive(Debug, Clone)]
508pub struct GridSampleOptions {
509    /// Interpolation mode (bilinear, nearest, or bicubic).
510    pub mode: InterpolateMode,
511    /// Padding mode for out-of-bounds coordinates.
512    pub padding_mode: GridSamplePaddingMode,
513    /// If `true`, grid values of -1 and 1 correspond to the corner pixels.
514    /// If `false`, they correspond to the corner points of the corner pixels
515    /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates).
516    pub align_corners: bool,
517}
518
519impl Default for GridSampleOptions {
520    fn default() -> Self {
521        Self {
522            mode: InterpolateMode::Bilinear,
523            padding_mode: GridSamplePaddingMode::Zeros,
524            align_corners: false,
525        }
526    }
527}
528
529impl From<InterpolateMode> for GridSampleOptions {
530    fn from(value: InterpolateMode) -> Self {
531        GridSampleOptions::new(value)
532    }
533}
534
535impl GridSampleOptions {
536    /// Create new grid sample options with the given interpolation mode.
537    ///
538    /// Uses default values for padding_mode (Zeros) and align_corners (false).
539    pub fn new(mode: InterpolateMode) -> Self {
540        Self {
541            mode,
542            ..Default::default()
543        }
544    }
545
546    /// Set the padding mode.
547    pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self {
548        self.padding_mode = padding_mode;
549        self
550    }
551
552    /// Set align_corners.
553    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
554        self.align_corners = align_corners;
555        self
556    }
557}
558
559/// Padding mode for tensor pad operations.
560///
561/// Defines how values are filled when padding a tensor beyond its original boundaries.
562/// Padding can be applied to any dimension of a tensor.
563///
564/// # Modes
565///
566/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0)
567/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size)
568/// - [`Edge`](PadMode::Edge): Replicate boundary values
569#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)]
570pub enum PadMode {
571    /// Fill padded regions with a constant value.
572    ///
573    /// # Example
574    /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0:
575    /// Result: `[0, 0, 1, 2, 3]`
576    Constant(f32),
577
578    /// Reflect values at the boundary, excluding the edge value.
579    ///
580    /// Padding must be less than the dimension size (i.e., `padding < dim_size`).
581    ///
582    /// # Example
583    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
584    /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0)
585    Reflect,
586
587    /// Replicate the edge values.
588    ///
589    /// # Example
590    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
591    /// Result: `[1, 1, 1, 2, 3, 4]`
592    Edge,
593}
594
595impl Default for PadMode {
596    fn default() -> Self {
597        PadMode::Constant(0.0)
598    }
599}
600
601impl<E: ElementConversion> From<E> for PadMode {
602    fn from(value: E) -> Self {
603        PadMode::Constant(value.elem())
604    }
605}
606
607/// Options for the attention module.
608#[derive(Debug, Clone, Copy, Default, PartialEq, serde::Deserialize, serde::Serialize)]
609pub struct AttentionModuleOptions {
610    /// Custom scale factor applied to QK^T. When `None`, defaults to `1/sqrt(head_dim)`.
611    pub scale: Option<f64>,
612
613    /// Soft capping applied before softmax: `softcap * tanh(scores / softcap)`.
614    /// Used by Gemma-2 and similar models. Must be positive when set.
615    pub softcap: Option<f64>,
616
617    /// When `true`, applies causal (autoregressive) masking so that each query position
618    /// can only attend to key positions at or before it. This is more efficient than
619    /// passing an explicit lower-triangular bool mask because backends can use optimized
620    /// kernel paths (e.g. flash attention with causal mode).
621    pub is_causal: bool,
622}
623
624/// Computation to be used to update the existing values in indexed assignment operations (scatter/select).
625#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
626pub enum IndexingUpdateOp {
627    /// Overwrite existing values.
628    Assign,
629    /// Performs an addition.
630    Add,
631    /// Multiply existing values.
632    Mul,
633    /// Take element-wise minimum.
634    Min,
635    /// Take element-wise maximum.
636    Max,
637}
638
639#[cfg(test)]
640mod tests {
641    use super::*;
642
643    #[test]
644    fn conv_options_symmetric_constructor() {
645        let options = ConvOptions::new([1, 2], [3, 4], [5, 6], 7);
646
647        assert_eq!(options.padding, [(3, 3), (4, 4)]);
648        assert_eq!(options.padding_begin(), [3, 4]);
649        assert_eq!(options.padding_end(), [3, 4]);
650        assert_eq!(options.symmetric_padding(), [3, 4]);
651        assert!(!options.is_asymmetric());
652    }
653
654    #[test]
655    fn conv_options_explicit_padding_constructor() {
656        let options = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
657
658        assert_eq!(options.padding, [(3, 4), (5, 6)]);
659        assert_eq!(options.padding_begin(), [3, 5]);
660        assert_eq!(options.padding_end(), [4, 6]);
661        assert!(options.is_asymmetric());
662    }
663
664    #[test]
665    #[allow(deprecated)]
666    fn padded_conv_options_convert_to_conv_options() {
667        let options: ConvOptions<2> =
668            PaddedConvOptions::asymmetric([1, 2], [3, 5], [4, 6], [7, 8], 9).into();
669
670        assert_eq!(options.padding, [(3, 4), (5, 6)]);
671        assert_eq!(options.stride, [1, 2]);
672        assert_eq!(options.dilation, [7, 8]);
673        assert_eq!(options.groups, 9);
674    }
675
676    #[test]
677    #[allow(deprecated)]
678    fn asymmetric_conv_options_roundtrip_through_padded_options() {
679        let expected = ConvOptions::new_with_padding([1, 2], [(3, 4), (5, 6)], [7, 8], 9);
680        let padded: PaddedConvOptions<2> = expected.clone().into();
681        let actual: ConvOptions<2> = padded.into();
682
683        assert_eq!(actual, expected);
684    }
685
686    #[test]
687    #[should_panic = "expected symmetric convolution padding"]
688    fn conv_options_symmetric_padding_with_asymmetric_options() {
689        let options = ConvOptions::new_with_padding([1], [(1, 2)], [1], 1);
690        let _ = options.symmetric_padding();
691    }
692
693    #[test]
694    #[should_panic = "stride must be non-zero"]
695    fn conv_options_stride_zero() {
696        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
697    }
698
699    #[test]
700    #[should_panic = "dilation must be non-zero"]
701    fn conv_options_dilation_zero() {
702        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
703    }
704
705    #[test]
706    #[should_panic = "groups must be non-zero"]
707    fn conv_options_groups_zero() {
708        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
709    }
710
711    #[test]
712    fn pool_options_defaults() {
713        let max = MaxPoolOptions::new([2, 3]);
714        let avg = AvgPoolOptions::new([2, 3]);
715
716        assert_eq!(max.stride, [2, 3]);
717        assert_eq!(max.padding, [(0, 0), (0, 0)]);
718        assert_eq!(max.dilation, [1, 1]);
719        assert!(!max.ceil_mode);
720        assert_eq!(avg.stride, [2, 3]);
721        assert_eq!(avg.padding, [(0, 0), (0, 0)]);
722        assert!(avg.count_include_pad);
723        assert!(!avg.ceil_mode);
724    }
725
726    #[test]
727    fn pool_options_padding() {
728        let symmetric = MaxPoolOptions::new([2, 2]).with_padding([0, 1]);
729        let asymmetric = AvgPoolOptions::new([2]).with_padding_pairs([(0, 1)]);
730
731        assert_eq!(symmetric.padding, [(0, 0), (1, 1)]);
732        assert!(!symmetric.is_asymmetric());
733        assert!(asymmetric.is_asymmetric());
734    }
735
736    #[test]
737    #[should_panic = "kernel size must be non-zero"]
738    fn max_pool_options_kernel_size_zero() {
739        let _opt = MaxPoolOptions::new([0, 1]);
740    }
741
742    #[test]
743    #[should_panic = "dilation must be non-zero"]
744    fn max_pool_options_dilation_zero() {
745        let _opt = MaxPoolOptions::new([1, 1]).with_dilation([0, 1]);
746    }
747
748    #[test]
749    #[should_panic = "stride must be non-zero"]
750    fn avg_pool_options_stride_zero() {
751        let _opt = AvgPoolOptions::new([1, 1]).with_stride([0, 1]);
752    }
753
754    #[test]
755    #[should_panic = "stride must be non-zero"]
756    fn conv_transpose_options_stride_zero() {
757        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
758    }
759
760    #[test]
761    #[should_panic = "dilation must be non-zero"]
762    fn conv_transpose_options_dilation_zero() {
763        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
764    }
765
766    #[test]
767    #[should_panic = "groups must be non-zero"]
768    fn conv_transpose_options_groups_zero() {
769        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
770    }
771
772    #[test]
773    #[should_panic = "stride must be non-zero"]
774    fn deform_conv_options_stride_zero() {
775        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
776    }
777
778    #[test]
779    #[should_panic = "dilation must be non-zero"]
780    fn deform_conv_options_dilation_zero() {
781        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
782    }
783
784    #[test]
785    #[should_panic = "weight groups must be non-zero"]
786    fn deform_conv_options_weights_groups_zero() {
787        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
788    }
789
790    #[test]
791    #[should_panic = "offset groups must be non-zero"]
792    fn deform_conv_options_offset_groups_zero() {
793        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
794    }
795
796    #[test]
797    #[should_panic = "stride must be non-zero"]
798    fn unfold_options_stride_zero() {
799        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
800    }
801
802    #[test]
803    #[should_panic = "dilation must be non-zero"]
804    fn unfold_options_dilation_zero() {
805        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
806    }
807}
burn_std/ops.rs

burn_std/
ops.rs
