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(crate) fn max_pool3d_composed<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(crate) fn max_pool3d_with_indices_composed<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/// Maximum-pool native volumes using the backend's volume operation.
119pub fn max_pool3d<B: Backend>(input: Tensor<B, 5>, kernel_size: [usize; 3], stride: [usize; 3],
120    padding: [usize; 3], dilation: [usize; 3], ceil_mode: bool) -> Tensor<B, 5> {
121    Tensor::new(TensorPrimitive::Float(B::max_pool3d(input.primitive.tensor(),
122        kernel_size, stride, padding, dilation, ceil_mode)))
123}
124
125/// Native volume maxima and flattened I64 positions in the original input.
126/// Empty selections are `-1`; tie/NaN selection retains spatial-then-depth order.
127pub fn max_pool3d_with_indices<B: Backend>(input: Tensor<B, 5>, kernel_size: [usize; 3],
128    stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], ceil_mode: bool)
129    -> (Tensor<B, 5>, Tensor<B, 5, Int>) {
130    let [_, _, depth, height, width] = input.dims();
131    assert!(depth > 0 && height > 0 && width > 0, "volume pooling indices require non-empty spatial axes");
132    let volume = depth.checked_mul(height).and_then(|size| size.checked_mul(width))
133        .expect("pooling volume size overflow");
134    assert!(volume <= i64::MAX as usize, "volume pooling positions exceed I64");
135    let output = B::max_pool3d_with_indices(input.primitive.tensor(), kernel_size, stride,
136        padding, dilation, ceil_mode);
137    (Tensor::new(TensorPrimitive::Float(output.output)), Tensor::new(output.indices))
138}
139
140/// Average-pool native `[batch, channels, depth, height, width]` activations.
141///
142/// Dispatches native volume pooling, retaining backend padding and ceil-window
143/// semantics. Backends without a specialized volume kernel use native spatial
144/// and depth pooling; no host reduction is used.
145pub fn avg_pool3d<B: Backend>(
146    input: Tensor<B, 5>,
147    kernel_size: [usize; 3],
148    stride: [usize; 3],
149    padding: [usize; 3],
150    count_include_pad: bool,
151    ceil_mode: bool,
152) -> Tensor<B, 5> {
153    Tensor::new(TensorPrimitive::Float(B::avg_pool3d(input.primitive.tensor(), kernel_size,
154        stride, padding, count_include_pad, ceil_mode)))
155}
156
157/// Adaptive average pooling to explicit `[depth, height, width]` extents.
158///
159/// Native adaptive pooling bins are retained, including overlapping bins and
160/// output extents larger than the input. Dispatches the backend's volume pooling
161/// operation; backends without a specialized kernel retain native separable pooling.
162pub fn adaptive_avg_pool3d<B: Backend>(
163    input: Tensor<B, 5>,
164    output_size: [usize; 3],
165) -> Tensor<B, 5> {
166    Tensor::new(TensorPrimitive::Float(B::adaptive_avg_pool3d(input.primitive.tensor(), output_size)))
167}
168
169fn average_excluding_explicit_padding<B: Backend, const D: usize, const N: usize>(
170    input: Tensor<B, D>,
171    padding: [(usize, usize); N],
172    pool: impl Fn(Tensor<B, D>) -> Tensor<B, D>,
173) -> Tensor<B, D> {
174    let storage = input.dtype();
175    let compute = if storage == DType::F64 { DType::F64 } else { DType::F32 };
176    let mut visible_shape = input.dims();
177    visible_shape[0] = 1;
178    visible_shape[1] = 1;
179    let visible = Tensor::<B, D>::ones(visible_shape, (&input.device(), compute))
180        .pad(padding, PadMode::Constant(0.0));
181    let values = pool(input.cast(compute).pad(padding, PadMode::Constant(0.0)));
182    let coverage = pool(visible);
183    (values / coverage).cast(storage)
184}
185
186/// Average pooling with explicit `(left, right)` padding.
187///
188/// When padding is excluded, the denominator counts original input elements,
189/// not the zeros materialized for asymmetric padding. Symmetric calls retain
190/// the original backend operation and arithmetic.
191pub fn avg_pool1d_padded<B: Backend>(
192    input: Tensor<B, 3>,
193    kernel_size: usize,
194    stride: usize,
195    padding: [(usize, usize); 1],
196    count_include_pad: bool,
197    ceil_mode: bool,
198) -> Tensor<B, 3> {
199    let [(left, right)] = padding;
200    if left == right {
201        return avg_pool1d(input, kernel_size, stride, left, count_include_pad, ceil_mode);
202    }
203    if count_include_pad {
204        return avg_pool1d(input.pad(padding, PadMode::Constant(0.0)),
205            kernel_size, stride, 0, true, ceil_mode);
206    }
207    average_excluding_explicit_padding(input, padding,
208        |input| avg_pool1d(input, kernel_size, stride, 0, true, ceil_mode))
209}
210
211/// Average pooling with explicit height and width `(before, after)` pairs.
212///
213/// Half-storage exclusion statistics use FP32; F64 inputs retain F64 statistics.
214/// The output retains the input storage, including its native gradient path.
215pub fn avg_pool2d_padded<B: Backend>(
216    input: Tensor<B, 4>,
217    kernel_size: [usize; 2],
218    stride: [usize; 2],
219    padding: [(usize, usize); 2],
220    count_include_pad: bool,
221    ceil_mode: bool,
222) -> Tensor<B, 4> {
223    if padding.iter().all(|(start, end)| start == end) {
224        return avg_pool2d(input, kernel_size, stride, padding.map(|(start, _)| start),
225            count_include_pad, ceil_mode);
226    }
227    if count_include_pad {
228        return avg_pool2d(input.pad(padding, PadMode::Constant(0.0)),
229            kernel_size, stride, [0; 2], true, ceil_mode);
230    }
231    average_excluding_explicit_padding(input, padding,
232        |input| avg_pool2d(input, kernel_size, stride, [0; 2], true, ceil_mode))
233}
234
235/// Average pooling with depth, height and width `(before, after)` pairs.
236///
237/// Coverage statistics broadcast across batch/channels and stay on the device.
238pub fn avg_pool3d_padded<B: Backend>(
239    input: Tensor<B, 5>,
240    kernel_size: [usize; 3],
241    stride: [usize; 3],
242    padding: [(usize, usize); 3],
243    count_include_pad: bool,
244    ceil_mode: bool,
245) -> Tensor<B, 5> {
246    if padding.iter().all(|(start, end)| start == end) {
247        return avg_pool3d(input, kernel_size, stride, padding.map(|(start, _)| start),
248            count_include_pad, ceil_mode);
249    }
250    if count_include_pad {
251        return avg_pool3d(input.pad(padding, PadMode::Constant(0.0)),
252            kernel_size, stride, [0; 3], true, ceil_mode);
253    }
254    average_excluding_explicit_padding(input, padding,
255        |input| avg_pool3d(input, kernel_size, stride, [0; 3], true, ceil_mode))
256}
257
258/// Maximum pooling with explicit `(left, right)` padding and native dilation.
259pub fn max_pool1d_padded<B: Backend>(
260    input: Tensor<B, 3>,
261    kernel_size: usize,
262    stride: usize,
263    padding: [(usize, usize); 1],
264    dilation: usize,
265    ceil_mode: bool,
266) -> Tensor<B, 3> {
267    let [(left, right)] = padding;
268    let (input, padding) = if left == right {
269        (input, left)
270    } else {
271        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), 0)
272    };
273    max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)
274}
275
276/// Maximum pooling with explicit height and width `(before, after)` pairs.
277pub fn max_pool2d_padded<B: Backend>(
278    input: Tensor<B, 4>,
279    kernel_size: [usize; 2],
280    stride: [usize; 2],
281    padding: [(usize, usize); 2],
282    dilation: [usize; 2],
283    ceil_mode: bool,
284) -> Tensor<B, 4> {
285    let (input, padding) = if padding.iter().all(|(start, end)| start == end) {
286        (input, padding.map(|(start, _)| start))
287    } else {
288        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), [0; 2])
289    };
290    max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
291}
292
293/// Maximum pooling with explicit depth, height and width padding pairs.
294pub fn max_pool3d_padded<B: Backend>(
295    input: Tensor<B, 5>,
296    kernel_size: [usize; 3],
297    stride: [usize; 3],
298    padding: [(usize, usize); 3],
299    dilation: [usize; 3],
300    ceil_mode: bool,
301) -> Tensor<B, 5> {
302    let (input, padding) = if padding.iter().all(|(start, end)| start == end) {
303        (input, padding.map(|(start, _)| start))
304    } else {
305        (input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)), [0; 3])
306    };
307    max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode)
308}
309
310fn unpad_pool_indices<B: Backend, const D: usize, const N: usize>(
311    indices: Tensor<B, D, Int>,
312    input_size: [usize; N],
313    padding: [(usize, usize); N],
314) -> Tensor<B, D, Int> {
315    let padded_size: [usize; N] = core::array::from_fn(|axis| {
316        input_size[axis].checked_add(padding[axis].0)
317            .and_then(|size| size.checked_add(padding[axis].1))
318            .expect("padded pooling index extent overflow")
319    });
320    let padded_volume = padded_size.iter().try_fold(1usize,
321        |size, axis| size.checked_mul(*axis)).expect("padded pooling index volume overflow");
322    let input_volume = input_size.iter().try_fold(1usize,
323        |size, axis| size.checked_mul(*axis)).expect("pooling input index volume overflow");
324    assert!(padded_volume > 0 && padded_volume <= i64::MAX as usize
325        && input_volume <= i64::MAX as usize, "pooling indices cannot be represented in I64");
326    let indices = indices.cast(DType::I64);
327    if input_volume == 0 {
328        return indices.zeros_like().sub_scalar(1);
329    }
330    let mut invalid = indices.clone().lower_elem(0)
331        .bool_or(indices.clone().greater_equal_elem(padded_volume as i64));
332    let mut remaining = indices.clone().clamp(0, padded_volume as i64 - 1);
333    let mut unpadded = indices.zeros_like();
334    let mut input_stride = 1usize;
335    for axis in (0..N).rev() {
336        let coordinate = remaining.clone().remainder_scalar(padded_size[axis] as i64)
337            .sub_scalar(padding[axis].0 as i64);
338        remaining = remaining.div_scalar(padded_size[axis] as i64);
339        invalid = invalid.bool_or(coordinate.clone().lower_elem(0))
340            .bool_or(coordinate.clone().greater_equal_elem(input_size[axis] as i64));
341        unpadded = unpadded + coordinate.clamp(0, input_size[axis].saturating_sub(1) as i64)
342            .mul_scalar(input_stride as i64);
343        input_stride = input_stride.checked_mul(input_size[axis])
344            .expect("pooling input index stride overflow");
345    }
346    unpadded.mask_fill(invalid, -1)
347}
348
349/// Maximum pooling with asymmetric length padding and original input positions.
350///
351/// Native symmetric indices retain their backend storage. Asymmetric indices
352/// are remapped on the device to I64 input coordinates; padding selections and
353/// native invalid indices are `-1`. Values retain the native gradient graph.
354pub fn max_pool1d_with_indices_padded<B: Backend>(
355    input: Tensor<B, 3>,
356    kernel_size: usize,
357    stride: usize,
358    padding: [(usize, usize); 1],
359    dilation: usize,
360    ceil_mode: bool,
361) -> (Tensor<B, 3>, Tensor<B, 3, Int>) {
362    let [(left, right)] = padding;
363    if left == right {
364        return max_pool1d_with_indices(input, kernel_size, stride, left, dilation, ceil_mode);
365    }
366    let [_, _, length] = input.dims();
367    let (values, indices) = max_pool1d_with_indices(
368        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
369        kernel_size, stride, 0, dilation, ceil_mode,
370    );
371    (values, unpad_pool_indices(indices, [length], padding))
372}
373
374/// Maximum pooling with asymmetric spatial padding and original `h * W + w` indices.
375///
376/// Indices exclude materialized padding, using `-1` for padding selections.
377pub fn max_pool2d_with_indices_padded<B: Backend>(
378    input: Tensor<B, 4>,
379    kernel_size: [usize; 2],
380    stride: [usize; 2],
381    padding: [(usize, usize); 2],
382    dilation: [usize; 2],
383    ceil_mode: bool,
384) -> (Tensor<B, 4>, Tensor<B, 4, Int>) {
385    if padding.iter().all(|(start, end)| start == end) {
386        return max_pool2d_with_indices(input, kernel_size, stride,
387            padding.map(|(start, _)| start), dilation, ceil_mode);
388    }
389    let [_, _, height, width] = input.dims();
390    let (values, indices) = max_pool2d_with_indices(
391        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
392        kernel_size, stride, [0; 2], dilation, ceil_mode,
393    );
394    (values, unpad_pool_indices(indices, [height, width], padding))
395}
396
397/// Maximum volume pooling with asymmetric padding and original flattened positions.
398///
399/// Returned I64 indices address the unpadded depth/height/width volume. Invalid
400/// native indices and selected materialized padding are `-1`.
401pub fn max_pool3d_with_indices_padded<B: Backend>(
402    input: Tensor<B, 5>,
403    kernel_size: [usize; 3],
404    stride: [usize; 3],
405    padding: [(usize, usize); 3],
406    dilation: [usize; 3],
407    ceil_mode: bool,
408) -> (Tensor<B, 5>, Tensor<B, 5, Int>) {
409    if padding.iter().all(|(start, end)| start == end) {
410        return max_pool3d_with_indices(input, kernel_size, stride,
411            padding.map(|(start, _)| start), dilation, ceil_mode);
412    }
413    let [_, _, depth, height, width] = input.dims();
414    let (values, indices) = max_pool3d_with_indices(
415        input.pad(padding, PadMode::Constant(f32::NEG_INFINITY)),
416        kernel_size, stride, [0; 3], dilation, ceil_mode,
417    );
418    (values, unpad_pool_indices(indices, [depth, height, width], padding))
419}
ruda_tensor/api/spatial_pool.rs

ruda_tensor/api/
spatial_pool.rs
