teeny_core/nn/
pool.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use core::marker::PhantomData;
18
19use crate::{
20    dtype::{Dtype, EagerTensor, Tensor},
21    nn::Layer,
22};
23
24// ---------------------------------------------------------------------------
25// Macro — reduces boilerplate for the nine pool variants
26// ---------------------------------------------------------------------------
27
28macro_rules! pool_layer_1d {
29    (#[doc = $doc:expr] $name:ident) => {
30        #[doc = $doc]
31        pub struct $name<D: Dtype, IT, OT, const RANK: usize> {
32            /// The pooling window length.
33            pub kernel_l: usize,
34            /// The stride between pooling windows.
35            pub stride: usize,
36            _pd: PhantomData<(D, IT, OT)>,
37        }
38
39        impl<D: Dtype, IT, OT, const RANK: usize> $name<D, IT, OT, RANK> {
40            /// Creates a new layer with the given kernel size and stride.
41            pub fn new(kernel_size: usize, stride: usize) -> Self {
42                Self {
43                    kernel_l: kernel_size,
44                    stride,
45                    _pd: PhantomData,
46                }
47            }
48        }
49
50        impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize>
51            Layer<IT> for $name<D, IT, OT, RANK>
52        {
53            type Output = OT;
54            fn call(&self, _input: IT) -> Self::Output {
55                todo!()
56            }
57        }
58    };
59}
60
61macro_rules! pool_layer_2d {
62    (#[doc = $doc:expr] $name:ident) => {
63        #[doc = $doc]
64        pub struct $name<D: Dtype, IT, OT, const RANK: usize> {
65            /// Pooling window height.
66            pub kernel_h: usize,
67            /// Pooling window width.
68            pub kernel_w: usize,
69            /// Vertical stride between pooling windows.
70            pub stride_h: usize,
71            /// Horizontal stride between pooling windows.
72            pub stride_w: usize,
73            /// Vertical padding applied before pooling.
74            pub padding_h: usize,
75            /// Horizontal padding applied before pooling.
76            pub padding_w: usize,
77            _pd: PhantomData<(D, IT, OT)>,
78        }
79
80        impl<D: Dtype, IT, OT, const RANK: usize> $name<D, IT, OT, RANK> {
81            /// Creates a new layer with the given kernel size and stride, no padding.
82            pub fn new(kernel_size: (usize, usize), stride: (usize, usize)) -> Self {
83                Self {
84                    kernel_h: kernel_size.0,
85                    kernel_w: kernel_size.1,
86                    stride_h: stride.0,
87                    stride_w: stride.1,
88                    padding_h: 0,
89                    padding_w: 0,
90                    _pd: PhantomData,
91                }
92            }
93
94            /// Creates a new layer with the given kernel size, stride, and padding.
95            pub fn with_padding(
96                kernel_size: (usize, usize),
97                stride: (usize, usize),
98                padding: (usize, usize),
99            ) -> Self {
100                Self {
101                    kernel_h: kernel_size.0,
102                    kernel_w: kernel_size.1,
103                    stride_h: stride.0,
104                    stride_w: stride.1,
105                    padding_h: padding.0,
106                    padding_w: padding.1,
107                    _pd: PhantomData,
108                }
109            }
110        }
111
112        impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize>
113            Layer<IT> for $name<D, IT, OT, RANK>
114        {
115            type Output = OT;
116            fn call(&self, _input: IT) -> Self::Output {
117                todo!()
118            }
119        }
120    };
121}
122
123macro_rules! pool_layer_3d {
124    (#[doc = $doc:expr] $name:ident) => {
125        #[doc = $doc]
126        pub struct $name<D: Dtype, IT, OT, const RANK: usize> {
127            /// Pooling window depth.
128            pub kernel_d: usize,
129            /// Pooling window height.
130            pub kernel_h: usize,
131            /// Pooling window width.
132            pub kernel_w: usize,
133            /// Stride along the depth dimension.
134            pub stride_d: usize,
135            /// Stride along the height dimension.
136            pub stride_h: usize,
137            /// Stride along the width dimension.
138            pub stride_w: usize,
139            _pd: PhantomData<(D, IT, OT)>,
140        }
141
142        impl<D: Dtype, IT, OT, const RANK: usize> $name<D, IT, OT, RANK> {
143            /// Creates a new layer with the given kernel size and stride.
144            pub fn new(kernel_size: (usize, usize, usize), stride: (usize, usize, usize)) -> Self {
145                Self {
146                    kernel_d: kernel_size.0,
147                    kernel_h: kernel_size.1,
148                    kernel_w: kernel_size.2,
149                    stride_d: stride.0,
150                    stride_h: stride.1,
151                    stride_w: stride.2,
152                    _pd: PhantomData,
153                }
154            }
155        }
156
157        impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize>
158            Layer<IT> for $name<D, IT, OT, RANK>
159        {
160            type Output = OT;
161            fn call(&self, _input: IT) -> Self::Output {
162                todo!()
163            }
164        }
165    };
166}
167
168// ---------------------------------------------------------------------------
169// AvgPool
170// ---------------------------------------------------------------------------
171
172pool_layer_1d!(
173    #[doc = "1-D average pooling."]
174    AvgPool1d
175);
176pool_layer_2d!(
177    #[doc = "2-D average pooling."]
178    AvgPool2d
179);
180pool_layer_3d!(
181    #[doc = "3-D average pooling."]
182    AvgPool3d
183);
184
185// ---------------------------------------------------------------------------
186// MaxPool
187// ---------------------------------------------------------------------------
188
189pool_layer_1d!(
190    #[doc = "1-D max pooling."]
191    MaxPool1d
192);
193pool_layer_2d!(
194    #[doc = "2-D max pooling."]
195    MaxPool2d
196);
197pool_layer_3d!(
198    #[doc = "3-D max pooling."]
199    MaxPool3d
200);
201
202// ---------------------------------------------------------------------------
203// LpPool — like AvgPool but with a configurable p-norm
204// ---------------------------------------------------------------------------
205
206/// 1-D power-average (Lp) pooling: like average pooling, generalized to an arbitrary p-norm.
207pub struct LpPool1d<D: Dtype, IT, OT, const RANK: usize> {
208    /// The pooling window length.
209    pub kernel_l: usize,
210    /// The stride between pooling windows.
211    pub stride: usize,
212    /// The `p` in the p-norm.
213    pub p: f64,
214    _pd: PhantomData<(D, IT, OT)>,
215}
216
217impl<D: Dtype, IT, OT, const RANK: usize> LpPool1d<D, IT, OT, RANK> {
218    /// Creates a new layer with the given kernel size, stride, and `p`-norm.
219    pub fn new(kernel_size: usize, stride: usize, p: f64) -> Self {
220        Self {
221            kernel_l: kernel_size,
222            stride,
223            p,
224            _pd: PhantomData,
225        }
226    }
227}
228
229impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
230    for LpPool1d<D, IT, OT, RANK>
231{
232    type Output = OT;
233    fn call(&self, _input: IT) -> Self::Output {
234        todo!()
235    }
236}
237
238/// 2-D power-average (Lp) pooling: like average pooling, generalized to an arbitrary p-norm.
239pub struct LpPool2d<D: Dtype, IT, OT, const RANK: usize> {
240    /// Pooling window height.
241    pub kernel_h: usize,
242    /// Pooling window width.
243    pub kernel_w: usize,
244    /// Vertical stride.
245    pub stride_h: usize,
246    /// Horizontal stride.
247    pub stride_w: usize,
248    /// The `p` in the p-norm.
249    pub p: f64,
250    _pd: PhantomData<(D, IT, OT)>,
251}
252
253impl<D: Dtype, IT, OT, const RANK: usize> LpPool2d<D, IT, OT, RANK> {
254    /// Creates a new layer with the given kernel size, stride, and `p`-norm.
255    pub fn new(kernel_size: (usize, usize), stride: (usize, usize), p: f64) -> Self {
256        Self {
257            kernel_h: kernel_size.0,
258            kernel_w: kernel_size.1,
259            stride_h: stride.0,
260            stride_w: stride.1,
261            p,
262            _pd: PhantomData,
263        }
264    }
265}
266
267impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
268    for LpPool2d<D, IT, OT, RANK>
269{
270    type Output = OT;
271    fn call(&self, _input: IT) -> Self::Output {
272        todo!()
273    }
274}
275
276/// 3-D power-average (Lp) pooling: like average pooling, generalized to an arbitrary p-norm.
277pub struct LpPool3d<D: Dtype, IT, OT, const RANK: usize> {
278    /// Pooling window depth.
279    pub kernel_d: usize,
280    /// Pooling window height.
281    pub kernel_h: usize,
282    /// Pooling window width.
283    pub kernel_w: usize,
284    /// Stride along the depth dimension.
285    pub stride_d: usize,
286    /// Stride along the height dimension.
287    pub stride_h: usize,
288    /// Stride along the width dimension.
289    pub stride_w: usize,
290    /// The `p` in the p-norm.
291    pub p: f64,
292    _pd: PhantomData<(D, IT, OT)>,
293}
294
295impl<D: Dtype, IT, OT, const RANK: usize> LpPool3d<D, IT, OT, RANK> {
296    /// Creates a new layer with the given kernel size, stride, and `p`-norm.
297    pub fn new(kernel_size: (usize, usize, usize), stride: (usize, usize, usize), p: f64) -> Self {
298        Self {
299            kernel_d: kernel_size.0,
300            kernel_h: kernel_size.1,
301            kernel_w: kernel_size.2,
302            stride_d: stride.0,
303            stride_h: stride.1,
304            stride_w: stride.2,
305            p,
306            _pd: PhantomData,
307        }
308    }
309}
310
311impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
312    for LpPool3d<D, IT, OT, RANK>
313{
314    type Output = OT;
315    fn call(&self, _input: IT) -> Self::Output {
316        todo!()
317    }
318}
teeny_core/nn/pool.rs

teeny_core/nn/
pool.rs
