mlx_rs/nn/
pooling.rs

1use std::iter::{once, zip};
2
3use crate::{error::Exception, module::Module, ops::as_strided, Array};
4use dyn_clone::DynClone;
5use mlx_macros::ModuleParameters;
6
7use crate::utils::SingleOrPair;
8
9/// Marker trait for pooling operations.
10pub trait Pooling
11where
12    Self: Fn(&Array, &[i32]) -> Result<Array, Exception> + DynClone,
13{
14}
15
16impl<T> Pooling for T where T: Fn(&Array, &[i32]) -> Result<Array, Exception> + DynClone {}
17
18/// Abstract pooling layer.
19///
20/// See also:
21///
22/// - [`MaxPool1d`]
23/// - [`MaxPool2d`]
24/// - [`AvgPool1d`]
25/// - [`AvgPool2d`]
26#[derive(ModuleParameters)]
27#[module(root = crate)]
28pub struct Pool {
29    /// Size of the pooling window
30    kernel_size: Vec<i32>,
31
32    /// Stride of the pooling window
33    stride: Vec<i64>,
34
35    /// Axes to pool over
36    axes: Vec<i32>,
37
38    /// Pooling operation
39    ///
40    /// TODO: We have Arc here just to make it `Clone` and `Send`. Is this necessary?
41    pooling_op: Box<dyn Pooling>,
42}
43
44impl Clone for Pool {
45    fn clone(&self) -> Self {
46        Self {
47            kernel_size: self.kernel_size.clone(),
48            stride: self.stride.clone(),
49            axes: self.axes.clone(),
50            pooling_op: dyn_clone::clone_box(&*self.pooling_op),
51        }
52    }
53}
54
55impl std::fmt::Debug for Pool {
56    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
57        f.debug_struct("Pool")
58            .field("kernel_size", &self.kernel_size)
59            .field("stride", &self.stride)
60            .field("axes", &self.axes)
61            .finish()
62    }
63}
64
65impl Pool {
66    /// Create a new abstract pooling layer.
67    pub fn new(kernel_size: Vec<i32>, stride: Vec<i64>, op: impl Pooling + 'static) -> Self {
68        let start = -(kernel_size.len() as i32) - 1;
69        let axes: Vec<_> = (start..-1).collect();
70        Self {
71            kernel_size,
72            stride,
73            axes,
74            pooling_op: Box::new(op),
75        }
76    }
77}
78
79impl Module<&Array> for Pool {
80    type Error = Exception;
81    type Output = Array;
82
83    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
84        let shape = x.shape();
85        let rest = &shape[1..shape.len() - 1];
86
87        let iter = zip(zip(rest, &self.kernel_size), &self.stride)
88            .map(|((size, window), stride)| (size - window) / *stride as i32 + 1);
89
90        let final_shape = once(shape[0])
91            .chain(iter)
92            .chain(self.kernel_size.iter().copied())
93            .chain(once(shape[shape.len() - 1]))
94            .collect::<Vec<_>>();
95
96        let strides = shape
97            .iter()
98            .map(|s| *s as i64)
99            .chain(once(1))
100            .rev()
101            .fold(vec![], |mut acc, a| {
102                match acc.last() {
103                    Some(&element) => acc.push(a * element),
104                    None => acc.push(a),
105                }
106                acc
107            })
108            .into_iter()
109            .rev()
110            .skip(1)
111            .collect::<Vec<_>>();
112        let middle_strides = &strides[1..strides.len() - 1];
113
114        let final_strides = once(strides[0])
115            .chain(zip(middle_strides, &self.stride).map(|(ms, s)| ms * s))
116            .chain(middle_strides.iter().copied())
117            .chain(once(1))
118            .collect::<Vec<_>>();
119
120        // TODO: double check if as_strided would ever panic
121        let strided = as_strided(x, &final_shape, &final_strides, None)?;
122        (self.pooling_op)(&strided, &self.axes)
123    }
124
125    fn training_mode(&mut self, _mode: bool) {}
126}
127
128macro_rules! impl_module {
129    ($name:ident) => {
130        impl Module<&Array> for $name {
131            type Output = Array;
132            type Error = Exception;
133
134            fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
135                self.inner.forward(x)
136            }
137
138            fn training_mode(&mut self, mode: bool) {
139                self.inner.training_mode(mode);
140            }
141        }
142    };
143}
144
145/// Applies 1-dimensional max pooling.
146///
147/// The input is expected to be `NLC`. The output will have the same N/C dimensions with the new `L
148/// = floor((L - kernel)/stride) + 1`
149///
150/// See [MaxPool1d python
151/// docs](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MaxPool1d.html)
152/// for more information.
153#[derive(Debug, Clone, ModuleParameters)]
154#[module(root = crate)]
155pub struct MaxPool1d {
156    #[param]
157    inner: Pool,
158}
159
160impl MaxPool1d {
161    /// Create a new 1-dimensional max pooling layer.
162    ///
163    /// # Params
164    ///
165    /// - `kernel_size`: The size of the pooling window.
166    /// - `stride`: The stride of the pooling window.
167    pub fn new(kernel_size: i32, stride: i64) -> Self {
168        let op = |x: &Array, axes: &[i32]| x.max_axes(axes, None);
169        let inner = Pool::new(vec![kernel_size], vec![stride], op);
170        Self { inner }
171    }
172}
173
174impl_module!(MaxPool1d);
175
176/// Applies 2-dimensional max pooling.
177///
178/// The input is expected to be `NHWC`. The output will have the same N/C dimensions with the new
179/// `H/W = floor((H/W - kernel)/stride) + 1`
180///
181/// See [MaxPool2d python
182/// docs](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MaxPool2d.html)
183/// for more information.
184#[derive(Debug, Clone, ModuleParameters)]
185#[module(root = crate)]
186pub struct MaxPool2d {
187    #[param]
188    inner: Pool,
189}
190
191impl MaxPool2d {
192    /// Create a new 2-dimensional max pooling layer.
193    ///
194    /// # Params
195    ///
196    /// - `kernel_size`: The size of the pooling window.
197    /// - `stride`: The stride of the pooling window.
198    pub fn new(
199        kernel_size: impl Into<SingleOrPair<i32>>,
200        stride: impl Into<SingleOrPair<i64>>,
201    ) -> Self {
202        let kernel_size = kernel_size.into();
203        let kernel_size = vec![kernel_size.first(), kernel_size.second()];
204        let stride = stride.into();
205        let stride = vec![stride.first(), stride.second()];
206
207        let op = |x: &Array, axes: &[i32]| x.max_axes(axes, None);
208        let inner = Pool::new(kernel_size, stride, op);
209        Self { inner }
210    }
211}
212
213impl_module!(MaxPool2d);
214
215/// Applies 1-dimensional average pooling.
216///
217/// The input is expected to be `NLC`. The output will have the same N/C dimensions with the new `L =
218/// floor((L - kernel)/stride) + 1`
219///
220/// See [AvgPool2d python
221/// docs](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.AvgPool2d.html)
222/// for more information.
223#[derive(Debug, Clone, ModuleParameters)]
224#[module(root = crate)]
225pub struct AvgPool1d {
226    #[param]
227    inner: Pool,
228}
229
230impl AvgPool1d {
231    /// Create a new 1-dimensional average pooling layer.
232    ///
233    /// # Params
234    ///
235    /// - `kernel_size`: The size of the pooling window.
236    /// - `stride`: The stride of the pooling window.
237    pub fn new(kernel_size: i32, stride: i64) -> Self {
238        let op = |x: &Array, axes: &[i32]| x.mean_axes(axes, None);
239        let inner = Pool::new(vec![kernel_size], vec![stride], op);
240        Self { inner }
241    }
242}
243
244impl_module!(AvgPool1d);
245
246/// Applies 2-dimensional average pooling.
247///
248/// The input is expected to be `NHWC`. The output will have the same N/C dimensions with the new
249/// `H/W = floor((H/W - kernel)/stride) + 1`
250///
251/// See [AvgPool2d python
252/// docs](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.AvgPool2d.html)
253/// for more information.
254#[derive(Debug, Clone, ModuleParameters)]
255#[module(root = crate)]
256pub struct AvgPool2d {
257    #[param]
258    inner: Pool,
259}
260
261impl AvgPool2d {
262    /// Create a new 2-dimensional average pooling layer.
263    ///
264    /// # Params
265    ///
266    /// - `kernel_size`: The size of the pooling window.
267    /// - `stride`: The stride of the pooling window.
268    pub fn new(
269        kernel_size: impl Into<SingleOrPair<i32>>,
270        stride: impl Into<SingleOrPair<i64>>,
271    ) -> Self {
272        let kernel_size = kernel_size.into();
273        let kernel_size = vec![kernel_size.first(), kernel_size.second()];
274        let stride = stride.into();
275        let stride = vec![stride.first(), stride.second()];
276
277        let op = |x: &Array, axes: &[i32]| x.mean_axes(axes, None);
278        let inner = Pool::new(kernel_size, stride, op);
279        Self { inner }
280    }
281}
282
283impl_module!(AvgPool2d);
284
285#[cfg(test)]
286mod tests {
287    use crate::{
288        array,
289        module::ModuleParameters,
290        test_utils::{assert_array_eq, tolerances},
291    };
292
293    use super::*;
294
295    #[test]
296    fn test_pool_has_no_learnable_params() {
297        let pool = MaxPool1d::new(2, 1);
298        let params = pool.parameters().flatten();
299        assert_eq!(params.len(), 0);
300    }
301
302    #[test]
303    fn test_max_pooling_1d_stride_1() {
304        let input = Array::from_iter(0..4, &[1, 4, 1]);
305        let mut pool = MaxPool1d::new(2, 1);
306        let output = pool.forward(&input).unwrap();
307        assert_array_eq(
308            output,
309            array!([1, 2, 3], shape = [1, 3, 1]),
310            tolerances::EXACT.rtol,
311            tolerances::EXACT.atol,
312        );
313    }
314
315    #[test]
316    fn test_max_pooling_1d_stride_2() {
317        let input = Array::from_iter(0..8, &[2, 4, 1]);
318        let mut pool = MaxPool1d::new(2, 2);
319        let output = pool.forward(&input).unwrap();
320        assert_array_eq(
321            output,
322            array!([1, 3, 5, 7], shape = [2, 2, 1]),
323            tolerances::EXACT.rtol,
324            tolerances::EXACT.atol,
325        );
326    }
327
328    #[test]
329    fn test_max_pooling_2d_stride_1() {
330        let input = Array::from_iter(0..16, &[1, 4, 4, 1]);
331        let mut pool = MaxPool2d::new(2, 1);
332        let output = pool.forward(&input).unwrap();
333        assert_array_eq(
334            output,
335            array!([5, 6, 7, 9, 10, 11, 13, 14, 15], shape = [1, 3, 3, 1]),
336            tolerances::EXACT.rtol,
337            tolerances::EXACT.atol,
338        );
339    }
340
341    #[test]
342    fn test_max_pooling_2d_stride_2() {
343        let input = Array::from_iter(0..32, &[2, 4, 4, 1]);
344        let mut pool = MaxPool2d::new(2, 2);
345        let output = pool.forward(&input).unwrap();
346        assert_array_eq(
347            output,
348            array!([5, 7, 13, 15, 21, 23, 29, 31], shape = [2, 2, 2, 1]),
349            tolerances::EXACT.rtol,
350            tolerances::EXACT.atol,
351        );
352    }
353
354    #[test]
355    fn test_avg_pooling_1d_stride_1() {
356        let input = Array::from_iter(0..4, &[1, 4, 1]);
357        let mut pool = AvgPool1d::new(2, 1);
358        let output = pool.forward(&input).unwrap();
359        assert_array_eq(
360            output,
361            array!([0.5, 1.5, 2.5], shape = [1, 3, 1]),
362            tolerances::EXACT.rtol,
363            tolerances::EXACT.atol,
364        );
365    }
366
367    #[test]
368    fn test_avg_pooling_1d_stride_2() {
369        let input = Array::from_iter(0..8, &[2, 4, 1]);
370        let mut pool = AvgPool1d::new(2, 2);
371        let output = pool.forward(&input).unwrap();
372        assert_array_eq(
373            output,
374            array!([0.5, 2.5, 4.5, 6.5], shape = [2, 2, 1]),
375            tolerances::EXACT.rtol,
376            tolerances::EXACT.atol,
377        );
378    }
379
380    #[test]
381    fn test_avg_pooling_2d_stride_1() {
382        let input = Array::from_iter(0..16, &[1, 4, 4, 1]);
383        let mut pool = AvgPool2d::new(2, 1);
384        let output = pool.forward(&input).unwrap();
385        assert_array_eq(
386            output,
387            array!(
388                [2.5, 3.5, 4.5, 6.5, 7.5, 8.5, 10.5, 11.5, 12.5],
389                shape = [1, 3, 3, 1]
390            ),
391            tolerances::EXACT.rtol,
392            tolerances::EXACT.atol,
393        );
394    }
395
396    #[test]
397    fn test_avg_pooling_2d_stride_2() {
398        let input = Array::from_iter(0..16, &[1, 4, 4, 1]);
399        let mut pool = AvgPool2d::new(2, 2);
400        let output = pool.forward(&input).unwrap();
401        assert_array_eq(
402            output,
403            array!([2.5, 4.5, 10.5, 12.5], shape = [1, 2, 2, 1]),
404            tolerances::EXACT.rtol,
405            tolerances::EXACT.atol,
406        );
407    }
408}
mlx_rs/nn/pooling.rs

mlx_rs/nn/
pooling.rs
