ruda_tensor/api/
spatial_pool.rs

1use super::{BasicOps, DType, Int, Tensor, TensorPrimitive, backend::Backend};
2use super::ops::PadMode;
3use super::module::{
4    avg_pool1d, avg_pool2d,
5    max_pool1d, max_pool2d, max_pool1d_with_indices, max_pool2d_with_indices,
6};
7
8pub(super) fn volume_planes<B: Backend, K: BasicOps<B>>(input: Tensor<B, 5, K>) -> Tensor<B, 4, K> {
9    let [batch, channels, depth, height, width] = input.dims();
10    let planes = batch.checked_mul(depth).expect("pooling plane count overflow");
11    input.permute([0, 2, 1, 3, 4]).reshape([planes, channels, height, width])
12}
13
14pub(super) fn plane_depth_lines<B: Backend, K: BasicOps<B>>(
15    planes: Tensor<B, 4, K>,
16    batch: usize,
17    depth: usize,
18) -> Tensor<B, 3, K> {
19    let [_, channels, height, width] = planes.dims();
20    let lines = batch.checked_mul(channels)
21        .and_then(|count| count.checked_mul(height))
22        .and_then(|count| count.checked_mul(width))
23        .expect("pooling depth line count overflow");
24    planes.reshape([batch, depth, channels, height, width])
25        .permute([0, 2, 3, 4, 1]).reshape([lines, 1, depth])
26}
27
28pub(super) fn depth_lines_volume<B: Backend, K: BasicOps<B>>(
29    lines: Tensor<B, 3, K>,
30    batch: usize,
31    channels: usize,
32    height: usize,
33    width: usize,
34) -> Tensor<B, 5, K> {
35    let [_, _, depth] = lines.dims();
36    lines.reshape([batch, channels, height, width, depth]).permute([0, 1, 4, 2, 3])
37}
38
39/// Pool native `[batch, channels, depth, height, width]` activations by maximum.
40///
41/// Executes the existing spatial and depth pooling kernels on the tensor's
42/// backend. Padding, dilation and ceil mode are applied independently per axis.
43/// Layout transformations and both pooling stages remain differentiable.
44pub fn max_pool3d<B: Backend>(
45    input: Tensor<B, 5>,
46    kernel_size: [usize; 3],
47    stride: [usize; 3],
48    padding: [usize; 3],
49    dilation: [usize; 3],
50    ceil_mode: bool,
51) -> Tensor<B, 5> {
52    let [batch, channels, depth, _, _] = input.dims();
53    let planes = max_pool2d(
54        volume_planes(input),
55        [kernel_size[1], kernel_size[2]],
56        [stride[1], stride[2]],
57        [padding[1], padding[2]],
58        [dilation[1], dilation[2]],
59        ceil_mode,
60    );
61    let [_, _, height, width] = planes.dims();
62    let lines = max_pool1d(
63        plane_depth_lines(planes, batch, depth),
64        kernel_size[0], stride[0], padding[0], dilation[0], ceil_mode,
65    );
66    depth_lines_volume(lines, batch, channels, height, width)
67}
68
69/// Maximum volume pooling with original input positions in native I64 storage.
70///
71/// Indices flatten each input volume as `d * (H * W) + h * W + w`, independently
72/// for every batch/channel. Tie selection follows the
73/// underlying spatial kernel followed by the depth kernel. Invalid backend
74/// indices become `-1`; they are never used as out-of-bounds gather addresses.
75pub fn max_pool3d_with_indices<B: Backend>(
76    input: Tensor<B, 5>,
77    kernel_size: [usize; 3],
78    stride: [usize; 3],
79    padding: [usize; 3],
80    dilation: [usize; 3],
81    ceil_mode: bool,
82) -> (Tensor<B, 5>, Tensor<B, 5, Int>) {
83    let [batch, channels, depth, input_height, input_width] = input.dims();
84    assert!(depth > 0 && input_height > 0 && input_width > 0,
85        "volume pooling indices require non-empty spatial axes");
86    let area = input_height.checked_mul(input_width).expect("pooling plane area overflow");
87    let volume = depth.checked_mul(area).expect("pooling volume size overflow");
88    assert!(volume <= i64::MAX as usize, "volume pooling indices exceed I64");
89    let (planes, plane_indices) = max_pool2d_with_indices(
90        volume_planes(input),
91        [kernel_size[1], kernel_size[2]],
92        [stride[1], stride[2]],
93        [padding[1], padding[2]],
94        [dilation[1], dilation[2]],
95        ceil_mode,
96    );
97    let [_, _, height, width] = planes.dims();
98    let (lines, depth_indices) = max_pool1d_with_indices(
99        plane_depth_lines(planes, batch, depth),
100        kernel_size[0], stride[0], padding[0], dilation[0], ceil_mode,
101    );
102    let depth_indices = depth_indices.cast(DType::I64);
103    let safe_depth_indices = depth_indices.clone().clamp(0, depth as i64 - 1);
104    let spatial_indices = plane_depth_lines(plane_indices.cast(DType::I64), batch, depth)
105        .gather(2, safe_depth_indices.clone());
106    let invalid = depth_indices.clone().lower_elem(0)
107        .bool_or(depth_indices.clone().greater_equal_elem(depth as i64))
108        .bool_or(spatial_indices.clone().lower_elem(0))
109        .bool_or(spatial_indices.clone().greater_equal_elem(area as i64));
110    let indices = (safe_depth_indices.mul_scalar(area as i64)
111        + spatial_indices.clamp(0, area as i64 - 1)).mask_fill(invalid, -1);
112    (
113        depth_lines_volume(lines, batch, channels, height, width),
114        depth_lines_volume(indices, batch, channels, height, width),
115    )
116}
117
118/// Average-pool native `[batch, channels, depth, height, width]` activations.
119///
120/// The rectangular reduction is factored into native spatial and depth pooling.
121/// Each stage retains the input's floating storage and its backend's arithmetic.
122/// Padding divisors factor per axis, including partially covered ceil windows.
123pub fn avg_pool3d<B: Backend>(
124    input: Tensor<B, 5>,
125    kernel_size: [usize; 3],
126    stride: [usize; 3],
127    padding: [usize; 3],
128    count_include_pad: bool,
129    ceil_mode: bool,
130) -> Tensor<B, 5> {
131    let [batch, channels, depth, _, _] = input.dims();
132    let planes = avg_pool2d(
133        volume_planes(input),
134        [kernel_size[1], kernel_size[2]],
135        [stride[1], stride[2]],
136        [padding[1], padding[2]],
137        count_include_pad,
138        ceil_mode,
139    );
140    let [_, _, height, width] = planes.dims();
141    let lines = avg_pool1d(
142        plane_depth_lines(planes, batch, depth),
143        kernel_size[0], stride[0], padding[0], count_include_pad, ceil_mode,
144    );
145    depth_lines_volume(lines, batch, channels, height, width)
146}
147
148/// Adaptive average pooling to explicit `[depth, height, width]` extents.
149///
150/// Native adaptive pooling bins are retained, including overlapping bins and
151/// output extents larger than the input. Dispatches the backend's volume pooling
152/// operation; backends without a specialized kernel retain native separable pooling.
153pub fn adaptive_avg_pool3d<B: Backend>(
154    input: Tensor<B, 5>,
155    output_size: [usize; 3],
156) -> Tensor<B, 5> {
157    Tensor::new(TensorPrimitive::Float(B::adaptive_avg_pool3d(input.primitive.tensor(), output_size)))
158}
159
160fn average_excluding_explicit_padding<B: Backend, const D: usize, const N: usize>(
161    input: Tensor<B, D>,
162    padding: [(usize, usize); N],
163    pool: impl Fn(Tensor<B, D>) -> Tensor<B, D>,
164) -> Tensor<B, D> {
165    let storage = input.dtype();
166    let compute = if storage == DType::F64 { DType::F64 } else { DType::F32 };
167    let mut visible_shape = input.dims();
168    visible_shape[0] = 1;
169    visible_shape[1] = 1;
170    let visible = Tensor::<B, D>::ones(visible_shape, (&input.device(), compute))
171        .pad(padding, PadMode::Constant(0.0));
172    let values = pool(input.cast(compute).pad(padding, PadMode::Constant(0.0)));
173    let coverage = pool(visible);
174    (values / coverage).cast(storage)
175}
176
177/// Average pooling with explicit `(left, right)` padding.
178///
179/// When padding is excluded, the denominator counts original input elements,
180/// not the zeros materialized for asymmetric padding. Symmetric calls retain
181/// the original backend operation and arithmetic.
182pub fn avg_pool1d_padded<B: Backend>(
183    input: Tensor<B, 3>,
184    kernel_size: usize,
185    stride: usize,
186    padding: [(usize, usize); 1],
187    count_include_pad: bool,
188    ceil_mode: bool,
189) -> Tensor<B, 3> {
190    let [(left, right)] = padding;
191    if left == right {
192        return avg_pool1d(input, kernel_size, stride, left, count_include_pad, ceil_mode);
193    }
194    if count_include_pad {
195        return avg_pool1d(input.pad(padding, PadMode::Constant(0.0)),
196            kernel_size, stride, 0, true, ceil_mode);
197    }
198    average_excluding_explicit_padding(input, padding,
199        |input| avg_pool1d(input, kernel_size, stride, 0, true, ceil_mode))
200}
201
202/// Average pooling with explicit height and width `(before, after)` pairs.
203///
204/// Half-storage exclusion statistics use FP32; F64 inputs retain F64 statistics.
205/// The output retains the input storage, including its native gradient path.
206pub fn avg_pool2d_padded<B: Backend>(
207    input: Tensor<B, 4>,
208    kernel_size: [usize; 2],
209    stride: [usize; 2],
210    padding: [(usize, usize); 2],
211    count_include_pad: bool,
212    ceil_mode: bool,
213) -> Tensor<B, 4> {
214    if padding.iter().all(|(start, end)| start == end) {
215        return avg_pool2d(input, kernel_size, stride, padding.map(|(start, _)| start),
216            count_include_pad, ceil_mode);
217    }
218    if count_include_pad {
219        return avg_pool2d(input.pad(padding, PadMode::Constant(0.0)),
220            kernel_size, stride, [0; 2], true, ceil_mode);
221    }
222    average_excluding_explicit_padding(input, padding,
223        |input| avg_pool2d(input, kernel_size, stride, [0; 2], true, ceil_mode))
224}
225
226/// Average pooling with depth, height and width `(before, after)` pairs.
227///
228/// Coverage statistics broadcast across batch/channels and stay on the device.
229pub fn avg_pool3d_padded<B: Backend>(
230    input: Tensor<B, 5>,
231    kernel_size: [usize; 3],
232    stride: [usize; 3],
233    padding: [(usize, usize); 3],
234    count_include_pad: bool,
235    ceil_mode: bool,
236) -> Tensor<B, 5> {
237    if padding.iter().all(|(start, end)| start == end) {
238        return avg_pool3d(input, kernel_size, stride, padding.map(|(start, _)| start),
239            count_include_pad, ceil_mode);
240    }
241    if count_include_pad {
242        return avg_pool3d(input.pad(padding, PadMode::Constant(0.0)),
243            kernel_size, stride, [0; 3], true, ceil_mode);
244    }
245    average_excluding_explicit_padding(input, padding,
246        |input| avg_pool3d(input, kernel_size, stride, [0; 3], true, ceil_mode))
247}
248
249/// Maximum pooling with explicit `(left, right)` padding and native dilation.
250pub fn max_pool1d_padded<B: Backend>(
251    input: Tensor<B, 3>,
252    kernel_size: usize,
253    stride: usize,
254    padding: [(usize, usize); 1],
255    dilation: usize,
256    ceil_mode: bool,
257) -> Tensor<B, 3> {
258    let [(left, right)] = padding;
259    let (input, padding) = if left == right {
260        (input, left)
261    } else {
262        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), 0)
263    };
264    max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)
265}
266
267/// Maximum pooling with explicit height and width `(before, after)` pairs.
268pub fn max_pool2d_padded<B: Backend>(
269    input: Tensor<B, 4>,
270    kernel_size: [usize; 2],
271    stride: [usize; 2],
272    padding: [(usize, usize); 2],
273    dilation: [usize; 2],
274    ceil_mode: bool,
275) -> Tensor<B, 4> {
276    let (input, padding) = if padding.iter().all(|(start, end)| start == end) {
277        (input, padding.map(|(start, _)| start))
278    } else {
279        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), [0; 2])
280    };
281    max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
282}
283
284/// Maximum pooling with explicit depth, height and width padding pairs.
285pub fn max_pool3d_padded<B: Backend>(
286    input: Tensor<B, 5>,
287    kernel_size: [usize; 3],
288    stride: [usize; 3],
289    padding: [(usize, usize); 3],
290    dilation: [usize; 3],
291    ceil_mode: bool,
292) -> Tensor<B, 5> {
293    let (input, padding) = if padding.iter().all(|(start, end)| start == end) {
294        (input, padding.map(|(start, _)| start))
295    } else {
296        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), [0; 3])
297    };
298    max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode)
299}
300
301fn unpad_pool_indices<B: Backend, const D: usize, const N: usize>(
302    indices: Tensor<B, D, Int>,
303    input_size: [usize; N],
304    padding: [(usize, usize); N],
305) -> Tensor<B, D, Int> {
306    let padded_size: [usize; N] = core::array::from_fn(|axis| {
307        input_size[axis].checked_add(padding[axis].0)
308            .and_then(|size| size.checked_add(padding[axis].1))
309            .expect("padded pooling index extent overflow")
310    });
311    let padded_volume = padded_size.iter().try_fold(1usize,
312        |size, axis| size.checked_mul(*axis)).expect("padded pooling index volume overflow");
313    let input_volume = input_size.iter().try_fold(1usize,
314        |size, axis| size.checked_mul(*axis)).expect("pooling input index volume overflow");
315    assert!(padded_volume > 0 && padded_volume <= i64::MAX as usize
316        && input_volume <= i64::MAX as usize, "pooling indices cannot be represented in I64");
317    let indices = indices.cast(DType::I64);
318    if input_volume == 0 {
319        return indices.zeros_like().sub_scalar(1);
320    }
321    let mut invalid = indices.clone().lower_elem(0)
322        .bool_or(indices.clone().greater_equal_elem(padded_volume as i64));
323    let mut remaining = indices.clone().clamp(0, padded_volume as i64 - 1);
324    let mut unpadded = indices.zeros_like();
325    let mut input_stride = 1usize;
326    for axis in (0..N).rev() {
327        let coordinate = remaining.clone().remainder_scalar(padded_size[axis] as i64)
328            .sub_scalar(padding[axis].0 as i64);
329        remaining = remaining.div_scalar(padded_size[axis] as i64);
330        invalid = invalid.bool_or(coordinate.clone().lower_elem(0))
331            .bool_or(coordinate.clone().greater_equal_elem(input_size[axis] as i64));
332        unpadded = unpadded + coordinate.clamp(0, input_size[axis].saturating_sub(1) as i64)
333            .mul_scalar(input_stride as i64);
334        input_stride = input_stride.checked_mul(input_size[axis])
335            .expect("pooling input index stride overflow");
336    }
337    unpadded.mask_fill(invalid, -1)
338}
339
340/// Maximum pooling with asymmetric length padding and original input positions.
341///
342/// Native symmetric indices retain their backend storage. Asymmetric indices
343/// are remapped on the device to I64 input coordinates; padding selections and
344/// native invalid indices are `-1`. Values retain the native gradient graph.
345pub fn max_pool1d_with_indices_padded<B: Backend>(
346    input: Tensor<B, 3>,
347    kernel_size: usize,
348    stride: usize,
349    padding: [(usize, usize); 1],
350    dilation: usize,
351    ceil_mode: bool,
352) -> (Tensor<B, 3>, Tensor<B, 3, Int>) {
353    let [(left, right)] = padding;
354    if left == right {
355        return max_pool1d_with_indices(input, kernel_size, stride, left, dilation, ceil_mode);
356    }
357    let [_, _, length] = input.dims();
358    let (values, indices) = max_pool1d_with_indices(
359        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
360        kernel_size, stride, 0, dilation, ceil_mode,
361    );
362    (values, unpad_pool_indices(indices, [length], padding))
363}
364
365/// Maximum pooling with asymmetric spatial padding and original `h * W + w` indices.
366///
367/// Indices exclude materialized padding, using `-1` for padding selections.
368pub fn max_pool2d_with_indices_padded<B: Backend>(
369    input: Tensor<B, 4>,
370    kernel_size: [usize; 2],
371    stride: [usize; 2],
372    padding: [(usize, usize); 2],
373    dilation: [usize; 2],
374    ceil_mode: bool,
375) -> (Tensor<B, 4>, Tensor<B, 4, Int>) {
376    if padding.iter().all(|(start, end)| start == end) {
377        return max_pool2d_with_indices(input, kernel_size, stride,
378            padding.map(|(start, _)| start), dilation, ceil_mode);
379    }
380    let [_, _, height, width] = input.dims();
381    let (values, indices) = max_pool2d_with_indices(
382        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
383        kernel_size, stride, [0; 2], dilation, ceil_mode,
384    );
385    (values, unpad_pool_indices(indices, [height, width], padding))
386}
387
388/// Maximum volume pooling with asymmetric padding and original flattened positions.
389///
390/// Returned I64 indices address the unpadded depth/height/width volume. Invalid
391/// native indices and selected materialized padding are `-1`.
392pub fn max_pool3d_with_indices_padded<B: Backend>(
393    input: Tensor<B, 5>,
394    kernel_size: [usize; 3],
395    stride: [usize; 3],
396    padding: [(usize, usize); 3],
397    dilation: [usize; 3],
398    ceil_mode: bool,
399) -> (Tensor<B, 5>, Tensor<B, 5, Int>) {
400    if padding.iter().all(|(start, end)| start == end) {
401        return max_pool3d_with_indices(input, kernel_size, stride,
402            padding.map(|(start, _)| start), dilation, ceil_mode);
403    }
404    let [_, _, depth, height, width] = input.dims();
405    let (values, indices) = max_pool3d_with_indices(
406        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
407        kernel_size, stride, [0; 3], dilation, ceil_mode,
408    );
409    (values, unpad_pool_indices(indices, [depth, height, width], padding))
410}
ruda_tensor/api/spatial_pool.rs

ruda_tensor/api/
spatial_pool.rs
