burn_backend/backend/ops/modules/
grid_sample.rs

1use crate::{
2    Backend, TensorMetadata,
3    element::ElementConversion,
4    ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode},
5    tensor::FloatTensor,
6};
7use alloc::vec;
8use burn_std::{Shape, Slice};
9
10/// Reference implementation of grid_sample_2d that supports all options.
11///
12/// # Arguments
13///
14/// * `tensor` - The tensor being sampled from, must be contiguous with shape (N, C, H_in, W_in)
15/// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1].
16///   A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right
17/// * `options` - Grid sampling options
18///
19/// # Returns
20///
21/// A tensor with shape (N, C, H_out, W_out)
22pub fn float_grid_sample_2d_ref<B: Backend>(
23    tensor: FloatTensor<B>,
24    grid: FloatTensor<B>,
25    options: GridSampleOptions,
26) -> FloatTensor<B> {
27    match options.mode {
28        InterpolateMode::Bilinear => float_grid_sample_2d_bilinear::<B>(
29            tensor,
30            grid,
31            options.padding_mode,
32            options.align_corners,
33        ),
34        _ => todo!(
35            "Default implementation for grid_sample_2d with {:?} unimplemented",
36            options.mode
37        ),
38    }
39}
40
41/// Bilinear grid sampling implementation.
42fn float_grid_sample_2d_bilinear<B: Backend>(
43    tensor: FloatTensor<B>,
44    grid: FloatTensor<B>,
45    padding_mode: GridSamplePaddingMode,
46    align_corners: bool,
47) -> FloatTensor<B> {
48    let n = tensor.shape().dims[0];
49    let c = tensor.shape().dims[1];
50    let h_in = tensor.shape().dims[2];
51    let w_in = tensor.shape().dims[3];
52    let h_out = grid.shape().dims[1];
53    let w_out = grid.shape().dims[2];
54    let spatial_in = h_in * w_in;
55    let spatial_out = h_out * w_out;
56
57    // Separate x and y coordinates from grid
58    // shape: (N, H_out, W_out, 1)
59    let grid_x_slice = vec![
60        Slice::new(0, Some(n as isize), 1),
61        Slice::new(0, Some(h_out as isize), 1),
62        Slice::new(0, Some(w_out as isize), 1),
63        Slice::new(0, Some(1), 1),
64    ];
65    let grid_y_slice = vec![
66        Slice::new(0, Some(n as isize), 1),
67        Slice::new(0, Some(h_out as isize), 1),
68        Slice::new(0, Some(w_out as isize), 1),
69        Slice::new(1, Some(2), 1),
70    ];
71
72    let grid_x = B::float_slice(grid.clone(), &grid_x_slice);
73    let grid_x = B::float_reshape(grid_x, Shape::new([n, 1, h_out, w_out]));
74    let grid_y = B::float_slice(grid.clone(), &grid_y_slice);
75    let grid_y = B::float_reshape(grid_y, Shape::new([n, 1, h_out, w_out]));
76
77    // Convert normalized grid coordinates [-1, 1] to pixel coordinates
78    let w_in_f = w_in as f64;
79    let h_in_f = h_in as f64;
80
81    let (grid_x, grid_y) = if align_corners {
82        // align_corners=true: x_pixel = (x_norm + 1) * (width - 1) / 2
83        // Maps -1 to 0 and 1 to width - 1
84        let grid_x = B::float_add_scalar(grid_x, 1.0f32.elem());
85        let grid_x = B::float_mul_scalar(grid_x, ((w_in_f - 1.0) / 2.0).elem());
86
87        let grid_y = B::float_add_scalar(grid_y, 1.0f32.elem());
88        let grid_y = B::float_mul_scalar(grid_y, ((h_in_f - 1.0) / 2.0).elem());
89
90        (grid_x, grid_y)
91    } else {
92        // align_corners=false: x_pixel = (x_norm + 1) * width / 2 - 0.5
93        // Maps -1 to -0.5 and 1 to width - 0.5
94        let grid_x = B::float_add_scalar(grid_x, 1.0f32.elem());
95        let grid_x = B::float_mul_scalar(grid_x, (w_in_f / 2.0).elem());
96        let grid_x = B::float_sub_scalar(grid_x, 0.5f32.elem());
97
98        let grid_y = B::float_add_scalar(grid_y, 1.0f32.elem());
99        let grid_y = B::float_mul_scalar(grid_y, (h_in_f / 2.0).elem());
100        let grid_y = B::float_sub_scalar(grid_y, 0.5f32.elem());
101
102        (grid_x, grid_y)
103    };
104
105    // Apply padding mode to coordinates
106    let (grid_x, grid_y) = match padding_mode {
107        GridSamplePaddingMode::Border => {
108            // Clamp coordinates to valid range [0, size-1]
109            let grid_x = B::float_clamp(grid_x, 0.0f32.elem(), ((w_in - 1) as f32).elem());
110            let grid_y = B::float_clamp(grid_y, 0.0f32.elem(), ((h_in - 1) as f32).elem());
111            (grid_x, grid_y)
112        }
113        GridSamplePaddingMode::Reflection => {
114            // Reflect coordinates at boundaries
115            let grid_x = reflect_coordinates::<B>(grid_x, w_in_f, align_corners);
116            let grid_y = reflect_coordinates::<B>(grid_y, h_in_f, align_corners);
117            (grid_x, grid_y)
118        }
119        GridSamplePaddingMode::Zeros => {
120            // Keep coordinates as-is, we'll mask out-of-bounds later
121            (grid_x, grid_y)
122        }
123    };
124
125    // Get floor indices for the four corners
126    let grid_x_floored = B::float_floor(grid_x.clone());
127    let grid_y_floored = B::float_floor(grid_y.clone());
128
129    // Compute interpolation weights (fractional part)
130    let x_frac = B::float_sub(grid_x.clone(), grid_x_floored.clone());
131    let y_frac = B::float_sub(grid_y.clone(), grid_y_floored.clone());
132
133    // Convert to integer indices
134    let x0 = B::float_into_int(grid_x_floored.clone());
135    let y0 = B::float_into_int(grid_y_floored.clone());
136    let x1 = B::float_into_int(B::float_add_scalar(grid_x_floored, 1.0f32.elem()));
137    let y1 = B::float_into_int(B::float_add_scalar(grid_y_floored, 1.0f32.elem()));
138
139    // Create masks for out-of-bounds coordinates (only used for zeros padding)
140    let (mask_00, mask_01, mask_10, mask_11) = if padding_mode == GridSamplePaddingMode::Zeros {
141        let x0_valid = B::int_greater_equal_elem(x0.clone(), 0.elem());
142        let x0_valid = B::bool_and(
143            x0_valid,
144            B::int_lower_elem(x0.clone(), (w_in as i32).elem()),
145        );
146        let x1_valid = B::int_greater_equal_elem(x1.clone(), 0.elem());
147        let x1_valid = B::bool_and(
148            x1_valid,
149            B::int_lower_elem(x1.clone(), (w_in as i32).elem()),
150        );
151        let y0_valid = B::int_greater_equal_elem(y0.clone(), 0.elem());
152        let y0_valid = B::bool_and(
153            y0_valid,
154            B::int_lower_elem(y0.clone(), (h_in as i32).elem()),
155        );
156        let y1_valid = B::int_greater_equal_elem(y1.clone(), 0.elem());
157        let y1_valid = B::bool_and(
158            y1_valid,
159            B::int_lower_elem(y1.clone(), (h_in as i32).elem()),
160        );
161
162        (
163            Some(B::bool_and(x0_valid.clone(), y0_valid.clone())),
164            Some(B::bool_and(x0_valid.clone(), y1_valid.clone())),
165            Some(B::bool_and(x1_valid.clone(), y0_valid)),
166            Some(B::bool_and(x1_valid, y1_valid)),
167        )
168    } else {
169        (None, None, None, None)
170    };
171
172    // Clamp indices to valid range for gather
173    let x0_clamped = B::int_clamp(x0, 0.elem(), ((w_in - 1) as i32).elem());
174    let x1_clamped = B::int_clamp(x1, 0.elem(), ((w_in - 1) as i32).elem());
175    let y0_clamped = B::int_clamp(y0, 0.elem(), ((h_in - 1) as i32).elem());
176    let y1_clamped = B::int_clamp(y1, 0.elem(), ((h_in - 1) as i32).elem());
177
178    // Linear indices: idx = y * W_in + x
179    let w_in_scalar: i32 = w_in as i32;
180    let idx_00 = B::int_add(
181        B::int_mul_scalar(y0_clamped.clone(), w_in_scalar.elem()),
182        x0_clamped.clone(),
183    );
184    let idx_01 = B::int_add(
185        B::int_mul_scalar(y1_clamped.clone(), w_in_scalar.elem()),
186        x0_clamped,
187    );
188    let idx_10 = B::int_add(
189        B::int_mul_scalar(y0_clamped, w_in_scalar.elem()),
190        x1_clamped.clone(),
191    );
192    let idx_11 = B::int_add(
193        B::int_mul_scalar(y1_clamped, w_in_scalar.elem()),
194        x1_clamped,
195    );
196
197    // [N, 1, H_out, W_out] -> [N, 1, H_out * W_out]
198    let idx_00 = B::int_reshape(idx_00, Shape::new([n, 1, spatial_out]));
199    let idx_01 = B::int_reshape(idx_01, Shape::new([n, 1, spatial_out]));
200    let idx_10 = B::int_reshape(idx_10, Shape::new([n, 1, spatial_out]));
201    let idx_11 = B::int_reshape(idx_11, Shape::new([n, 1, spatial_out]));
202
203    // [N, 1, spatial] -> [N, C, spatial]
204    let idx_00 = B::int_expand(idx_00, Shape::new([n, c, spatial_out]));
205    let idx_01 = B::int_expand(idx_01, Shape::new([n, c, spatial_out]));
206    let idx_10 = B::int_expand(idx_10, Shape::new([n, c, spatial_out]));
207    let idx_11 = B::int_expand(idx_11, Shape::new([n, c, spatial_out]));
208
209    let tensor_flat = B::float_reshape(tensor, Shape::new([n, c, spatial_in]));
210
211    let sample_00 = B::float_gather(2, tensor_flat.clone(), idx_00);
212    let sample_01 = B::float_gather(2, tensor_flat.clone(), idx_01);
213    let sample_10 = B::float_gather(2, tensor_flat.clone(), idx_10);
214    let sample_11 = B::float_gather(2, tensor_flat, idx_11);
215
216    // Reshape samples to (N, C, H_out, W_out)
217    let sample_00 = B::float_reshape(sample_00, Shape::new([n, c, h_out, w_out]));
218    let sample_01 = B::float_reshape(sample_01, Shape::new([n, c, h_out, w_out]));
219    let sample_10 = B::float_reshape(sample_10, Shape::new([n, c, h_out, w_out]));
220    let sample_11 = B::float_reshape(sample_11, Shape::new([n, c, h_out, w_out]));
221
222    // Apply masks for zeros padding (set out-of-bounds samples to 0)
223    let (sample_00, sample_01, sample_10, sample_11) =
224        if padding_mode == GridSamplePaddingMode::Zeros {
225            let mask_00 = mask_00.unwrap();
226            let mask_01 = mask_01.unwrap();
227            let mask_10 = mask_10.unwrap();
228            let mask_11 = mask_11.unwrap();
229
230            let mask_00_inv = B::bool_not(mask_00);
231            let mask_00_inv = B::bool_reshape(mask_00_inv, Shape::new([n, 1, h_out, w_out]));
232            let mask_00_inv = B::bool_expand(mask_00_inv, Shape::new([n, c, h_out, w_out]));
233            let mask_01_inv = B::bool_not(mask_01);
234            let mask_01_inv = B::bool_reshape(mask_01_inv, Shape::new([n, 1, h_out, w_out]));
235            let mask_01_inv = B::bool_expand(mask_01_inv, Shape::new([n, c, h_out, w_out]));
236            let mask_10_inv = B::bool_not(mask_10);
237            let mask_10_inv = B::bool_reshape(mask_10_inv, Shape::new([n, 1, h_out, w_out]));
238            let mask_10_inv = B::bool_expand(mask_10_inv, Shape::new([n, c, h_out, w_out]));
239            let mask_11_inv = B::bool_not(mask_11);
240            let mask_11_inv = B::bool_reshape(mask_11_inv, Shape::new([n, 1, h_out, w_out]));
241            let mask_11_inv = B::bool_expand(mask_11_inv, Shape::new([n, c, h_out, w_out]));
242
243            (
244                B::float_mask_fill(sample_00, mask_00_inv, 0.0f32.elem()),
245                B::float_mask_fill(sample_01, mask_01_inv, 0.0f32.elem()),
246                B::float_mask_fill(sample_10, mask_10_inv, 0.0f32.elem()),
247                B::float_mask_fill(sample_11, mask_11_inv, 0.0f32.elem()),
248            )
249        } else {
250            (sample_00, sample_01, sample_10, sample_11)
251        };
252
253    // Compute bilinear interpolation weights
254    let one_minus_x = B::float_neg(x_frac.clone());
255    let one_minus_x = B::float_add_scalar(one_minus_x, 1.0f32.elem());
256
257    let one_minus_y = B::float_neg(y_frac.clone());
258    let one_minus_y = B::float_add_scalar(one_minus_y, 1.0f32.elem());
259
260    let weight_00 = B::float_mul(one_minus_x.clone(), one_minus_y.clone());
261    let weight_01 = B::float_mul(one_minus_x.clone(), y_frac.clone());
262    let weight_10 = B::float_mul(x_frac.clone(), one_minus_y);
263    let weight_11 = B::float_mul(x_frac, y_frac);
264
265    // Bilinear interpolation
266    let result = B::float_mul(sample_00, weight_00);
267    let result = B::float_add(result, B::float_mul(sample_01, weight_01));
268    let result = B::float_add(result, B::float_mul(sample_10, weight_10));
269
270    B::float_add(result, B::float_mul(sample_11, weight_11))
271}
272
273/// Reflect coordinates at boundaries using a triangle wave pattern.
274///
275/// For align_corners=true: reflects within [0, size-1]
276/// For align_corners=false: reflects within [-0.5, size-0.5]
277fn reflect_coordinates<B: Backend>(
278    coords: FloatTensor<B>,
279    size: f64,
280    align_corners: bool,
281) -> FloatTensor<B> {
282    let (min_val, max_val) = if align_corners {
283        (0.0f32, (size - 1.0) as f32)
284    } else {
285        (-0.5f32, (size - 0.5) as f32)
286    };
287
288    let span = max_val - min_val;
289    if span <= 0.0 {
290        // Edge case: size is 1, just return min_val everywhere
291        let zeros = B::float_mul_scalar(coords, 0.0f32.elem());
292        return B::float_add_scalar(zeros, min_val.elem());
293    }
294
295    // Triangle wave formula: span - |((x mod 2*span) - span)| + min_val
296    let period = 2.0 * span;
297
298    // x = abs(coord - min_val)
299    let x = B::float_sub_scalar(coords, min_val.elem());
300    let x = B::float_abs(x);
301
302    // x_mod = x - floor(x / period) * period
303    let x_div = B::float_div_scalar(x.clone(), period.elem());
304    let x_div_floor = B::float_floor(x_div);
305    let x_mod = B::float_sub(x, B::float_mul_scalar(x_div_floor, period.elem()));
306
307    // result = span - abs(x_mod - span) + min_val
308    let diff = B::float_sub_scalar(x_mod, span.elem());
309    let abs_diff = B::float_abs(diff);
310    let reflected = B::float_sub_scalar(abs_diff, span.elem());
311    let reflected = B::float_neg(reflected);
312    B::float_add_scalar(reflected, min_val.elem())
313}
burn_backend/backend/ops/modules/grid_sample.rs

burn_backend/backend/ops/modules/
grid_sample.rs
