cubek-pool 0.3.0

CubeK: Pool
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

use crate::definition::{
    AdaptiveAvgPoolOptions, AvgPoolOptions, MaxPoolOptions, PoolForwardProblem, PoolMode,
};
use cubecl::zspace::Shape;

pub trait PoolGeometry<const N: usize> {
    fn output_shape(&self, input_shape: &Shape) -> Shape;
}

impl<const N: usize> PoolGeometry<N> for PoolForwardProblem<N> {
    fn output_shape(&self, input_shape: &Shape) -> Shape {
        self.mode.output_shape(input_shape)
    }
}

impl<const N: usize> PoolGeometry<N> for PoolMode<N> {
    fn output_shape(&self, input_shape: &Shape) -> Shape {
        match self {
            PoolMode::Max(opts) => opts.output_shape(input_shape),
            PoolMode::Avg(opts) => opts.output_shape(input_shape),
            PoolMode::AdaptiveAvg(opts) => opts.output_shape(input_shape),
        }
    }
}

impl<const N: usize> PoolGeometry<N> for MaxPoolOptions<N> {
    fn output_shape(&self, input_shape: &Shape) -> Shape {
        let input_dims = input_shape.to_vec();
        let mut out = vec![input_dims[0]];
        for i in 0..N {
            let idx = 1 + i;
            let effective_kernel = (self.window.kernel_size[i] - 1) * self.dilation[i] + 1;
            let padded = input_dims[idx] + 2 * self.window.padding[i];
            out.push(pooled_dim(
                padded,
                effective_kernel,
                self.window.stride[i],
                self.window.ceil_mode,
            ));
        }
        out.push(input_dims[input_dims.len() - 1]);
        Shape::from(out)
    }
}

impl<const N: usize> PoolGeometry<N> for AvgPoolOptions<N> {
    fn output_shape(&self, input_shape: &Shape) -> Shape {
        let input_dims = input_shape.to_vec();
        let mut out = vec![input_dims[0]];
        for i in 0..N {
            let idx = 1 + i;
            let padded = input_dims[idx] + 2 * self.window.padding[i];
            let effective_kernel = self.window.kernel_size[i];
            out.push(pooled_dim(
                padded,
                effective_kernel,
                self.window.stride[i],
                self.window.ceil_mode,
            ));
        }
        out.push(input_dims[input_dims.len() - 1]);
        Shape::from(out)
    }
}

impl<const N: usize> PoolGeometry<N> for AdaptiveAvgPoolOptions<N> {
    fn output_shape(&self, _input_shape: &Shape) -> Shape {
        let input_dims = _input_shape.to_vec();
        let mut out = vec![input_dims[0]];
        out.extend_from_slice(&self.output_size);
        out.push(input_dims[input_dims.len() - 1]);
        Shape::from(out)
    }
}

fn pooled_dim(padded: usize, effective_kernel: usize, stride: usize, ceil_mode: bool) -> usize {
    let size = (padded as f32 - effective_kernel as f32) / stride as f32;
    (if ceil_mode { size.ceil() } else { size.floor() }) as usize + 1
}