hpt_traits/ops/
pooling.rs

1use hpt_common::{error::base::TensorError, shape::shape::Shape};
2
3/// trait for pooling that the output type is the same as the input type
4pub trait NormalPooling {
5    /// the output type is the same as the input type
6    type Output;
7
8    /// Performs a 2D max pooling operation on the input tensor, selecting the maximum value from each window.
9    ///
10    /// ## Parameters:
11    /// `kernels`: Shape of the pooling window, typically `[kernel_height, kernel_width]`
12    ///
13    /// `steps`: Stride of the pooling operation as `[step_height, step_width]`
14    ///
15    /// `padding`: Padding size as `[(padding_top, padding_bottom), (padding_left, padding_right)]`
16    ///
17    /// `dilation`: Kernel dilation factors as `[dilation_height, dilation_width]`
18    ///
19    /// ## Example:
20    /// ```rust
21    /// let input = Tensor::<f32>::randn([1, 32, 32, 16])?;
22    /// let output = input.maxpool2d(
23    ///     [2, 2],           // kernel size
24    ///     [2, 2],           // stride
25    ///     [(0, 0), (0, 0)], // padding
26    ///     [1, 1],           // dilation
27    /// )?; // shape: [1, 16, 16, 16]
28    /// ```
29    fn maxpool2d<S: Into<Shape>>(
30        &self,
31        kernels_shape: S,
32        steps: [i64; 2],
33        padding: [(i64, i64); 2],
34        dilation: [i64; 2],
35    ) -> Result<Self::Output, TensorError>;
36
37    /// Performs an adaptive max pooling operation on the input tensor, automatically determining the kernel size and stride to produce the specified output dimensions.
38    ///
39    /// ## Parameters:
40    /// `output_size`: Desired output spatial dimensions as `[out_height, out_width]`
41    ///
42    /// ## Example:
43    /// ```rust
44    /// let input = Tensor::<f32>::randn([1, 32, 32, 16])?;
45    /// let output = input.adaptive_maxpool2d([16, 16])?; // shape: [1, 16, 16, 16]
46    /// let output2 = input.adaptive_maxpool2d([8, 8])?; // shape: [1, 8, 8, 16]
47    /// ```
48    fn adaptive_maxpool2d(&self, output_size: [i64; 2]) -> Result<Self::Output, TensorError>;
49}
50
51/// trait for pooling that the output type is the same as the input type
52pub trait FloatOutPooling {
53    /// the output type is the same as the input type
54    type Output;
55
56    /// Performs a 2D average pooling operation on the input tensor, computing the average value from each window.
57    ///
58    /// ## Parameters:
59    /// `kernels`: Shape of the pooling window, typically `[kernel_height, kernel_width]`
60    ///
61    /// `steps`: Stride of the pooling operation as `[step_height, step_width]`
62    ///
63    /// `padding`: Padding size as `[(padding_top, padding_bottom), (padding_left, padding_right)]`
64    ///
65    /// `dilation`: Kernel dilation factors as `[dilation_height, dilation_width]`
66    ///
67    /// ## Example:
68    /// ```rust
69    /// let input = Tensor::<f32>::randn([1, 32, 32, 16])?;
70    /// let output = input.avgpool2d(
71    ///     [2, 2],           // kernel size
72    ///     [2, 2],           // stride
73    ///     [(0, 0), (0, 0)], // padding
74    ///     [1, 1],           // dilation
75    /// )?; // shape: [1, 16, 16, 16]
76    /// ```
77    fn avgpool2d<S: Into<Shape>>(
78        &self,
79        kernels_shape: S,
80        steps: [i64; 2],
81        padding: [(i64, i64); 2],
82        dilation: [i64; 2],
83    ) -> Result<Self::Output, TensorError>;
84
85    /// Performs an adaptive avg pooling operation on the input tensor, automatically determining the kernel size and stride to produce the specified output dimensions.
86    ///
87    /// ## Parameters:
88    /// `output_size`: Desired output spatial dimensions as `[out_height, out_width]`
89    ///
90    /// ## Example:
91    /// ```rust
92    /// let input = Tensor::<f32>::randn([1, 32, 32, 16])?;
93    /// let output = input.adaptive_avgpool2d([16, 16])?; // shape: [1, 16, 16, 16]
94    /// let output2 = input.adaptive_avgpool2d([8, 8])?; // shape: [1, 8, 8, 16]
95    /// ```
96    fn adaptive_avgpool2d(&self, output_size: [i64; 2]) -> Result<Self::Output, TensorError>;
97}
hpt_traits/ops/pooling.rs

hpt_traits/ops/
pooling.rs
