burn_tensor/tensor/api/
pad.rs

1use alloc::vec::Vec;
2
3use crate::{Tensor, kind::Numeric, ops::PadMode};
4
5/// Trait for types that can be used as padding specifications.
6///
7/// Padding is specified as `(before, after)` pairs per dimension, returned as a
8/// fixed-size array `[(usize, usize); D]`. If fewer pairs than dimensions are provided,
9/// they apply to the **last** N dimensions (earlier dimensions are left unpadded).
10pub trait IntoPadding<const D: usize> {
11    /// Converts into a fixed-size array of `(before, after)` padding pairs.
12    fn into_padding(self) -> [(usize, usize); D];
13}
14
15impl<const D: usize, const N: usize> IntoPadding<D> for [(usize, usize); N] {
16    fn into_padding(self) -> [(usize, usize); D] {
17        assert!(
18            N <= D,
19            "Padding has {} pairs but tensor only has {} dimensions",
20            N,
21            D
22        );
23        let mut result = [(0usize, 0usize); D];
24        let offset = D - N;
25        for (i, pair) in self.into_iter().enumerate() {
26            result[offset + i] = pair;
27        }
28        result
29    }
30}
31
32/// Backward-compatible: `(left, right, top, bottom)` maps to last 2 dimensions.
33///
34/// Equivalent to `[(top, bottom), (left, right)]`.
35impl<const D: usize> IntoPadding<D> for (usize, usize, usize, usize) {
36    fn into_padding(self) -> [(usize, usize); D] {
37        let (left, right, top, bottom) = self;
38        let mut result = [(0usize, 0usize); D];
39        result[D - 2] = (top, bottom);
40        result[D - 1] = (left, right);
41        result
42    }
43}
44
45impl<const D: usize> IntoPadding<D> for &[(usize, usize)] {
46    fn into_padding(self) -> [(usize, usize); D] {
47        assert!(
48            self.len() <= D,
49            "Padding has {} pairs but tensor only has {} dimensions",
50            self.len(),
51            D
52        );
53        let mut result = [(0usize, 0usize); D];
54        let offset = D - self.len();
55        for (i, &pair) in self.iter().enumerate() {
56            result[offset + i] = pair;
57        }
58        result
59    }
60}
61
62impl<const D: usize> IntoPadding<D> for Vec<(usize, usize)> {
63    fn into_padding(self) -> [(usize, usize); D] {
64        assert!(
65            self.len() <= D,
66            "Padding has {} pairs but tensor only has {} dimensions",
67            self.len(),
68            D
69        );
70        let mut result = [(0usize, 0usize); D];
71        let offset = D - self.len();
72        for (i, pair) in self.into_iter().enumerate() {
73            result[offset + i] = pair;
74        }
75        result
76    }
77}
78
79impl<const D: usize, K> Tensor<D, K>
80where
81    K: Numeric,
82{
83    /// Pads the tensor using the specified padding mode.
84    ///
85    /// Padding is specified as `(before, after)` pairs. If fewer pairs than tensor dimensions
86    /// are provided, they apply to the **last** N dimensions (unspecified leading dimensions
87    /// are left unpadded).
88    ///
89    /// For backward compatibility, a `(left, right, top, bottom)` tuple is also accepted,
90    /// which pads the last two dimensions.
91    ///
92    /// # Arguments
93    ///
94    /// * `padding` - Padding specification. Accepts:
95    ///   - `[(before, after); N]` fixed-size array of pairs (N <= D)
96    ///   - `&[(before, after)]` slice of pairs per dimension
97    ///   - `Vec<(before, after)>` vector of pairs
98    ///   - `(left, right, top, bottom)` tuple for last-2-dim backward compatibility
99    /// * `mode` - The padding mode: `Constant(value)`, `Reflect`, or `Edge`.
100    ///
101    /// # Returns
102    ///
103    /// A new tensor with the specified padding applied.
104    ///
105    /// # Panics
106    ///
107    /// - Panics if more padding pairs are provided than tensor dimensions.
108    /// - `Reflect` mode panics if padding exceeds `dimension_size - 1`.
109    /// - `Edge` mode panics if padding is applied to a zero-sized dimension.
110    ///
111    /// # Example
112    ///
113    /// ```rust
114    /// use burn_tensor::{Tensor, Shape};
115    /// use burn_tensor::ops::PadMode;
116    ///
117    /// let device = Default::default();
118    /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
119    ///
120    /// // Constant padding with value 0.0 (backward-compatible tuple)
121    /// let padded = tensor.clone().pad((1, 1, 1, 1), PadMode::Constant(0.0));
122    ///
123    /// // Pad arbitrary dimensions with slice of (before, after) pairs
124    /// let padded = tensor.clone().pad([(1, 1), (2, 2)], PadMode::Constant(0.0));
125    ///
126    /// // Pad only the last dimension
127    /// let padded = tensor.pad([(1, 1)], PadMode::Reflect);
128    /// ```
129    pub fn pad(self, padding: impl IntoPadding<D>, mode: impl Into<PadMode>) -> Self {
130        let pairs = padding.into_padding();
131        Tensor::new(K::pad(self.primitive, &pairs, mode.into()))
132    }
133}
burn_tensor/tensor/api/pad.rs

burn_tensor/tensor/api/
pad.rs
