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

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

#[derive(Clone, Debug)]
pub enum PoolProblem {
    Forward(PoolForward),
    Backward(PoolBackward),
}

#[derive(Clone, Debug)]
pub enum PoolForward {
    D1(PoolForwardProblem<1>),
    D2(PoolForwardProblem<2>),
    D3(PoolForwardProblem<3>),
}

#[derive(Clone, Debug)]
pub enum PoolBackward {
    D1(PoolBackwardProblem<1>),
    D2(PoolBackwardProblem<2>),
    D3(PoolBackwardProblem<3>),
}

#[derive(Clone, Debug)]
pub struct PoolForwardProblem<const N: usize> {
    pub input_shape: Shape,
    pub with_indices: bool,
    pub mode: PoolMode<N>,
}

#[derive(Clone, Debug)]
pub struct PoolBackwardProblem<const N: usize> {
    pub input_size: [usize; N],
    pub out_grad_shape: Shape,
    pub with_indices: bool,
    pub mode: PoolMode<N>,
}

#[cfg(feature = "benchmarks")]
impl<const N: usize> PoolBackwardProblem<N> {
    /// Reconstruct the channels-last input shape from the stored problem dimensions.
    pub(crate) fn input_shape(&self) -> Shape {
        let mut shape = vec![self.out_grad_shape[0]];
        shape.extend_from_slice(&self.input_size);
        shape.push(self.out_grad_shape[N + 1]);
        Shape::from(shape)
    }
}

#[derive(Clone, Debug)]
pub enum PoolMode<const N: usize> {
    Max(MaxPoolOptions<N>),
    Avg(AvgPoolOptions<N>),
    AdaptiveAvg(AdaptiveAvgPoolOptions<N>),
}

impl<const N: usize> From<MaxPoolOptions<N>> for PoolMode<N> {
    fn from(options: MaxPoolOptions<N>) -> Self {
        PoolMode::Max(options)
    }
}

impl<const N: usize> From<AvgPoolOptions<N>> for PoolMode<N> {
    fn from(options: AvgPoolOptions<N>) -> Self {
        PoolMode::Avg(options)
    }
}

impl<const N: usize> From<AdaptiveAvgPoolOptions<N>> for PoolMode<N> {
    fn from(options: AdaptiveAvgPoolOptions<N>) -> Self {
        PoolMode::AdaptiveAvg(options)
    }
}