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.
20    pub padding: [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,
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
46/// Convolution options with support for asymmetric padding.
47///
48/// Wraps [`ConvOptions`] (which represents symmetric padding for the backend op)
49/// and adds optional asymmetric padding. When asymmetric padding is specified,
50/// the functional convolution layer applies an explicit pad operation before
51/// dispatching to the backend.
52///
53/// Implements `From<ConvOptions<N>>` for backward compatibility.
54#[derive(Debug, Clone)]
55pub struct PaddedConvOptions<const N: usize> {
56    /// The underlying convolution options for the backend.
57    pub options: ConvOptions<N>,
58    /// Padding at the end of each dimension (e.g., bottom/right for 2D).
59    /// If `None`, padding is symmetric (same as `options.padding`).
60    /// If `Some`, specifies different end-padding per dimension.
61    pub padding_end: Option<[usize; N]>,
62}
63
64impl<const N: usize> PaddedConvOptions<N> {
65    /// Creates options with asymmetric padding.
66    ///
67    /// `padding_start` is stored in `ConvOptions::padding`.
68    /// `padding_end` specifies the end padding per dimension.
69    pub fn asymmetric(
70        stride: [usize; N],
71        padding_start: [usize; N],
72        padding_end: [usize; N],
73        dilation: [usize; N],
74        groups: usize,
75    ) -> Self {
76        let options = ConvOptions::new(stride, padding_start, dilation, groups);
77        if padding_start == padding_end {
78            Self {
79                options,
80                padding_end: None,
81            }
82        } else {
83            Self {
84                options,
85                padding_end: Some(padding_end),
86            }
87        }
88    }
89
90    /// Returns true if padding is asymmetric.
91    pub fn is_asymmetric(&self) -> bool {
92        self.padding_end.is_some()
93    }
94}
95
96impl<const N: usize> From<ConvOptions<N>> for PaddedConvOptions<N> {
97    fn from(options: ConvOptions<N>) -> Self {
98        Self {
99            options,
100            padding_end: None,
101        }
102    }
103}
104
105/// Deformable convolution options.
106#[derive(Debug, Clone, Hash, PartialEq, Eq)]
107pub struct DeformConvOptions<const N: usize> {
108    /// Stride (non-zero).
109    pub stride: [usize; N],
110
111    /// Padding.
112    pub padding: [usize; N],
113
114    /// Dilation (non-zero).
115    pub dilation: [usize; N],
116
117    /// Weight Groups (non-zero).
118    pub weight_groups: usize,
119
120    /// Offset Groups (non-zero).
121    pub offset_groups: usize,
122}
123
124impl<const N: usize> DeformConvOptions<N> {
125    /// Constructs a new `DeformConvOptions`.
126    pub fn new(
127        stride: [usize; N],
128        padding: [usize; N],
129        dilation: [usize; N],
130        weight_groups: usize,
131        offset_groups: usize,
132    ) -> Self {
133        Self {
134            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
135            padding,
136            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
137            weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"),
138            offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"),
139        }
140    }
141}
142
143/// Transposed convolution options.
144#[derive(Debug, Clone, Hash, PartialEq, Eq)]
145pub struct ConvTransposeOptions<const N: usize> {
146    /// Stride (non-zero).
147    pub stride: [usize; N],
148
149    /// Padding.
150    pub padding: [usize; N],
151
152    /// Padding out.
153    pub padding_out: [usize; N],
154
155    /// Dilation (non-zero).
156    pub dilation: [usize; N],
157
158    /// Groups (non-zero).
159    pub groups: usize,
160}
161
162impl<const N: usize> ConvTransposeOptions<N> {
163    /// Constructs a new `ConvTransposeOptions`.
164    pub fn new(
165        stride: [usize; N],
166        padding: [usize; N],
167        padding_out: [usize; N],
168        dilation: [usize; N],
169        groups: usize,
170    ) -> Self {
171        Self {
172            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
173            padding,
174            padding_out,
175            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
176            groups: check_nonzero(groups, "groups must be non-zero"),
177        }
178    }
179}
180
181/// Unfold operation options.
182#[derive(Debug, Clone)]
183pub struct UnfoldOptions {
184    /// The number of positions to slide over the input tensor in each dimension.
185    /// A stride of `[1, 1]` will slide the kernel one pixel at a time.
186    pub stride: [usize; 2],
187
188    /// The number of zero-padding pixels added to each side of the input tensor in each dimension.
189    pub padding: [usize; 2],
190
191    /// The spacing between the blocks (patches) in the original input tensor.
192    pub dilation: [usize; 2],
193}
194
195impl UnfoldOptions {
196    /// Constructs a new `UnfoldOptions`.
197    pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self {
198        Self {
199            stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")),
200            padding,
201            dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")),
202        }
203    }
204}
205
206/// Algorithm used.
207#[derive(new, Debug, Clone, serde::Deserialize, serde::Serialize)]
208pub enum InterpolateMode {
209    /// Nearest-neighbor floor interpolation.
210    /// Matches the legacy behavior of OpenCV’s INTER_NEAREST. It results in a bottom-right shift when resizing.
211    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
212    Nearest,
213
214    /// Nearest-neighbor exact interpolation.
215    /// <https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation>
216    NearestExact,
217
218    /// Bilinear interpolation.
219    /// <https://en.wikipedia.org/wiki/Bilinear_interpolation>
220    Bilinear,
221
222    /// Bicubic interpolation.
223    /// <https://en.wikipedia.org/wiki/Bicubic_interpolation>
224    Bicubic,
225
226    /// Lanczos3 interpolation (6-tap sinc-based filter).
227    /// <https://en.wikipedia.org/wiki/Lanczos_resampling>
228    Lanczos3,
229}
230
231/// Interpolation options.
232#[derive(Debug, Clone)]
233pub struct InterpolateOptions {
234    /// Algorithm used.
235    pub mode: InterpolateMode,
236    /// If `true`, the input and output tensors are aligned by their corner pixels.
237    /// If `false`, half-pixel coordinate mapping is used instead.
238    pub align_corners: bool,
239}
240
241impl InterpolateOptions {
242    /// Create new interpolate options with the given mode.
243    /// Defaults to `align_corners = true`.
244    pub fn new(mode: InterpolateMode) -> Self {
245        Self {
246            mode,
247            align_corners: true,
248        }
249    }
250
251    /// Set align_corners.
252    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
253        self.align_corners = align_corners;
254        self
255    }
256}
257
258/// Padding mode for grid sampling when coordinates are out of bounds.
259///
260/// Matches PyTorch's `padding_mode` parameter in `grid_sample`.
261#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)]
262pub enum GridSamplePaddingMode {
263    /// Fill with zeros for out-of-bounds coordinates.
264    #[default]
265    Zeros,
266    /// Clamp coordinates to the border (use nearest edge value).
267    Border,
268    /// Reflect coordinates at the boundary.
269    Reflection,
270}
271
272/// Options for grid sampling operations.
273#[derive(Debug, Clone)]
274pub struct GridSampleOptions {
275    /// Interpolation mode (bilinear, nearest, or bicubic).
276    pub mode: InterpolateMode,
277    /// Padding mode for out-of-bounds coordinates.
278    pub padding_mode: GridSamplePaddingMode,
279    /// If `true`, grid values of -1 and 1 correspond to the corner pixels.
280    /// If `false`, they correspond to the corner points of the corner pixels
281    /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates).
282    pub align_corners: bool,
283}
284
285impl Default for GridSampleOptions {
286    fn default() -> Self {
287        Self {
288            mode: InterpolateMode::Bilinear,
289            padding_mode: GridSamplePaddingMode::Zeros,
290            align_corners: false,
291        }
292    }
293}
294
295impl From<InterpolateMode> for GridSampleOptions {
296    fn from(value: InterpolateMode) -> Self {
297        GridSampleOptions::new(value)
298    }
299}
300
301impl GridSampleOptions {
302    /// Create new grid sample options with the given interpolation mode.
303    ///
304    /// Uses default values for padding_mode (Zeros) and align_corners (false).
305    pub fn new(mode: InterpolateMode) -> Self {
306        Self {
307            mode,
308            ..Default::default()
309        }
310    }
311
312    /// Set the padding mode.
313    pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self {
314        self.padding_mode = padding_mode;
315        self
316    }
317
318    /// Set align_corners.
319    pub fn with_align_corners(mut self, align_corners: bool) -> Self {
320        self.align_corners = align_corners;
321        self
322    }
323}
324
325/// Padding mode for tensor pad operations.
326///
327/// Defines how values are filled when padding a tensor beyond its original boundaries.
328/// Padding can be applied to any dimension of a tensor.
329///
330/// # Modes
331///
332/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0)
333/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size)
334/// - [`Edge`](PadMode::Edge): Replicate boundary values
335#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)]
336pub enum PadMode {
337    /// Fill padded regions with a constant value.
338    ///
339    /// # Example
340    /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0:
341    /// Result: `[0, 0, 1, 2, 3]`
342    Constant(f32),
343
344    /// Reflect values at the boundary, excluding the edge value.
345    ///
346    /// Padding must be less than the dimension size (i.e., `padding < dim_size`).
347    ///
348    /// # Example
349    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
350    /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0)
351    Reflect,
352
353    /// Replicate the edge values.
354    ///
355    /// # Example
356    /// For tensor `[1, 2, 3, 4]` with padding 2 on the left:
357    /// Result: `[1, 1, 1, 2, 3, 4]`
358    Edge,
359}
360
361impl Default for PadMode {
362    fn default() -> Self {
363        PadMode::Constant(0.0)
364    }
365}
366
367impl<E: ElementConversion> From<E> for PadMode {
368    fn from(value: E) -> Self {
369        PadMode::Constant(value.elem())
370    }
371}
372
373/// Options for the attention module.
374#[derive(Debug, Clone, Copy, Default, PartialEq, serde::Deserialize, serde::Serialize)]
375pub struct AttentionModuleOptions {
376    /// Custom scale factor applied to QK^T. When `None`, defaults to `1/sqrt(head_dim)`.
377    pub scale: Option<f64>,
378
379    /// Soft capping applied before softmax: `softcap * tanh(scores / softcap)`.
380    /// Used by Gemma-2 and similar models. Must be positive when set.
381    pub softcap: Option<f64>,
382
383    /// When `true`, applies causal (autoregressive) masking so that each query position
384    /// can only attend to key positions at or before it. This is more efficient than
385    /// passing an explicit lower-triangular bool mask because backends can use optimized
386    /// kernel paths (e.g. flash attention with causal mode).
387    pub is_causal: bool,
388}
389
390/// Computation to be used to update the existing values in indexed assignment operations (scatter/select).
391#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
392pub enum IndexingUpdateOp {
393    /// Overwrite existing values.
394    Assign,
395    /// Performs an addition.
396    Add,
397    /// Multiply existing values.
398    Mul,
399    /// Take element-wise minimum.
400    Min,
401    /// Take element-wise maximum.
402    Max,
403}
404
405#[cfg(test)]
406mod tests {
407    use super::*;
408
409    #[test]
410    #[should_panic = "stride must be non-zero"]
411    fn conv_options_stride_zero() {
412        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
413    }
414
415    #[test]
416    #[should_panic = "dilation must be non-zero"]
417    fn conv_options_dilation_zero() {
418        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
419    }
420
421    #[test]
422    #[should_panic = "groups must be non-zero"]
423    fn conv_options_groups_zero() {
424        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
425    }
426
427    #[test]
428    #[should_panic = "stride must be non-zero"]
429    fn conv_transpose_options_stride_zero() {
430        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
431    }
432
433    #[test]
434    #[should_panic = "dilation must be non-zero"]
435    fn conv_transpose_options_dilation_zero() {
436        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
437    }
438
439    #[test]
440    #[should_panic = "groups must be non-zero"]
441    fn conv_transpose_options_groups_zero() {
442        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
443    }
444
445    #[test]
446    #[should_panic = "stride must be non-zero"]
447    fn deform_conv_options_stride_zero() {
448        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
449    }
450
451    #[test]
452    #[should_panic = "dilation must be non-zero"]
453    fn deform_conv_options_dilation_zero() {
454        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
455    }
456
457    #[test]
458    #[should_panic = "weight groups must be non-zero"]
459    fn deform_conv_options_weights_groups_zero() {
460        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
461    }
462
463    #[test]
464    #[should_panic = "offset groups must be non-zero"]
465    fn deform_conv_options_offset_groups_zero() {
466        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
467    }
468
469    #[test]
470    #[should_panic = "stride must be non-zero"]
471    fn unfold_options_stride_zero() {
472        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
473    }
474
475    #[test]
476    #[should_panic = "dilation must be non-zero"]
477    fn unfold_options_dilation_zero() {
478        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
479    }
480}
burn_std/ops.rs

burn_std/
ops.rs
