mlx_rs/nn/
convolution.rs

1use crate::module::{Module, Param};
2use crate::{
3    error::{ConvBuildError, Exception},
4    ops::{conv1d, conv2d, zeros},
5    random::uniform,
6    Array,
7};
8use mlx_internal_macros::{Buildable, Builder};
9use mlx_macros::ModuleParameters;
10
11use crate::utils::{SingleOrPair, SingleOrTriple};
12
13/// Builder for the `Conv1d` module.
14#[derive(Debug, Clone, Builder)]
15#[builder(
16    root = crate,
17    build_with = build_conv1d,
18    err = ConvBuildError,
19)]
20pub struct Conv1dBuilder {
21    /// Number of input channels.
22    pub input_channels: i32,
23
24    /// Number of output channels.
25    pub output_channels: i32,
26
27    /// Size of the convolution filters.
28    pub kernel_size: i32,
29
30    /// If `true`, add a learnable bias to the output. Default to [`Conv1d::DEFAULT_BIAS`] if not
31    /// specified.
32    #[builder(optional, default = Conv1d::DEFAULT_BIAS)]
33    pub bias: bool,
34
35    /// Stride. Default to [`Conv1d::DEFAULT_STRIDE`] if not specified.
36    #[builder(optional, default = Conv1d::DEFAULT_STRIDE)]
37    pub stride: i32,
38
39    /// Padding. Default to [`Conv1d::DEFAULT_PADDING`] if not specified.
40    #[builder(optional, default = Conv1d::DEFAULT_PADDING)]
41    pub padding: i32,
42
43    /// Dilation. Default to [`Conv1d::DEFAULT_DILATION`] if not specified.
44    #[builder(optional, default = Conv1d::DEFAULT_DILATION)]
45    pub dilation: i32,
46
47    /// Groups. Default to [`Conv1d::DEFAULT_GROUPS`] if not specified.
48    #[builder(optional, default = Conv1d::DEFAULT_GROUPS)]
49    pub groups: i32,
50}
51
52fn build_conv1d(builder: Conv1dBuilder) -> Result<Conv1d, ConvBuildError> {
53    let input_channels = builder.input_channels;
54    let output_channels = builder.output_channels;
55    let kernel_size = builder.kernel_size;
56    let with_bias = builder.bias;
57    let groups = checked_groups(input_channels, output_channels, builder.groups)?;
58
59    let scale = f32::sqrt(1.0f32 / (input_channels * kernel_size) as f32);
60    let weight = uniform::<_, f32>(
61        -scale,
62        scale,
63        &[output_channels, kernel_size, input_channels / groups],
64        None,
65    )?;
66    let bias = if with_bias {
67        Some(zeros::<f32>(&[output_channels])?)
68    } else {
69        None
70    };
71
72    Ok(Conv1d {
73        weight: Param::new(weight),
74        bias: Param::new(bias),
75        stride: builder.stride,
76        padding: builder.padding,
77        dilation: builder.dilation,
78        groups,
79    })
80}
81
82/// Applies a 1-dimensional convolution over the multi-channel input sequence.
83///
84/// The channels are expected to be last i.e. the input shape should be `NLC` where:
85///
86/// - `N` is the batch dimension
87/// - `L` is the sequence length
88/// - `C` is the number of input channels
89#[derive(Debug, Clone, ModuleParameters, Buildable)]
90#[module(root = crate)]
91#[buildable(root = crate)]
92pub struct Conv1d {
93    /// The weight of the convolution layer.
94    #[param]
95    pub weight: Param<Array>,
96
97    /// The bias of the convolution layer.
98    #[param]
99    pub bias: Param<Option<Array>>,
100
101    /// Stride. Default to [`Conv1d::DEFAULT_STRIDE`] if not specified.
102    pub stride: i32,
103
104    /// Padding. Default to [`Conv1d::DEFAULT_PADDING`] if not specified.
105    pub padding: i32,
106
107    /// Dilation. Default to [`Conv1d::DEFAULT_DILATION`] if not specified.
108    pub dilation: i32,
109
110    /// Groups. Default to [`Conv1d::DEFAULT_GROUPS`] if not specified.
111    pub groups: i32,
112}
113
114impl Conv1d {
115    /// Default value for `with_bias` if not specified.
116    pub const DEFAULT_BIAS: bool = true;
117
118    /// Default value for `stride` if not specified.
119    pub const DEFAULT_STRIDE: i32 = 1;
120
121    /// Default value for `padding` if not specified.
122    pub const DEFAULT_PADDING: i32 = 0;
123
124    /// Default value for `dilation` if not specified.
125    pub const DEFAULT_DILATION: i32 = 1;
126
127    /// Default value for `groups` if not specified.
128    pub const DEFAULT_GROUPS: i32 = 1;
129}
130
131impl Module<&Array> for Conv1d {
132    type Error = Exception;
133    type Output = Array;
134
135    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
136        let mut y = conv1d(
137            x,
138            self.weight.as_ref(),
139            self.stride,
140            self.padding,
141            self.dilation,
142            self.groups,
143        )?;
144        if let Some(bias) = &self.bias.value {
145            y += bias;
146        }
147        Ok(y)
148    }
149
150    fn training_mode(&mut self, _: bool) {}
151}
152
153/// Builder for the `Conv2d` module.
154#[derive(Debug, Clone, Builder)]
155#[builder(
156    root = crate,
157    build_with = build_conv2d,
158    err = ConvBuildError,
159)]
160pub struct Conv2dBuilder {
161    /// Number of input channels.
162    pub input_channels: i32,
163
164    /// Number of output channels.
165    pub output_channels: i32,
166
167    /// Size of the convolution filters.
168    pub kernel_size: SingleOrPair<i32>,
169
170    /// If `true`, add a learnable bias to the output. Default to [`Conv2d::DEFAULT_BIAS`] if not
171    /// specified.
172    #[builder(optional, default = Conv2d::DEFAULT_BIAS)]
173    pub bias: bool,
174
175    /// Stride. Default to [`Conv2d::DEFAULT_STRIDE`] if not specified.
176    #[builder(optional, default = Conv2d::DEFAULT_STRIDE)]
177    pub stride: SingleOrPair<i32>,
178
179    /// Padding. Default to [`Conv2d::DEFAULT_PADDING`] if not specified.
180    #[builder(optional, default = Conv2d::DEFAULT_PADDING)]
181    pub padding: SingleOrPair<i32>,
182
183    /// Dilation. Default to [`Conv2d::DEFAULT_DILATION`] if not specified.
184    #[builder(optional, default = Conv2d::DEFAULT_DILATION)]
185    pub dilation: SingleOrPair<i32>,
186
187    /// Groups. Default to [`Conv2d::DEFAULT_GROUPS`] if not specified.
188    #[builder(optional, default = Conv2d::DEFAULT_GROUPS)]
189    pub groups: i32,
190}
191
192fn build_conv2d(builder: Conv2dBuilder) -> Result<Conv2d, ConvBuildError> {
193    let input_channels = builder.input_channels;
194    let output_channels = builder.output_channels;
195    let kernel_size: (i32, i32) = builder.kernel_size.into();
196    let with_bias = builder.bias;
197    let padding = builder.padding.into();
198    let stride = builder.stride.into();
199    let dilation = builder.dilation.into();
200    let groups = checked_groups(input_channels, output_channels, builder.groups)?;
201
202    let scale = f32::sqrt(1.0 / (input_channels * kernel_size.0 * kernel_size.1) as f32);
203    let weight = uniform::<_, f32>(
204        -scale,
205        scale,
206        &[
207            output_channels,
208            kernel_size.0,
209            kernel_size.1,
210            input_channels / groups,
211        ],
212        None,
213    )?;
214    let bias = if with_bias {
215        Some(zeros::<f32>(&[output_channels])?)
216    } else {
217        None
218    };
219
220    Ok(Conv2d {
221        weight: Param::new(weight),
222        bias: Param::new(bias),
223        stride,
224        padding,
225        dilation,
226        groups,
227    })
228}
229
230/// Applies a 2-dimensional convolution over the multi-channel input image.
231///
232/// The channels are expected to be last i.e. the input shape should be `NHWC` where:
233///
234/// - `N` is the batch dimension
235/// - `H` is the input image height
236/// - `W` is the input image width
237/// - `C` is the number of input channels
238#[derive(Debug, Clone, ModuleParameters, Buildable)]
239#[module(root = crate)]
240#[buildable(root = crate)]
241pub struct Conv2d {
242    /// The weight of the convolution layer.
243    #[param]
244    pub weight: Param<Array>,
245
246    /// The bias of the convolution layer.
247    #[param]
248    pub bias: Param<Option<Array>>,
249
250    /// Stride. Default to [`Conv2d::DEFAULT_STRIDE`] if not specified.
251    pub stride: (i32, i32),
252
253    /// Padding. Default to [`Conv2d::DEFAULT_PADDING`] if not specified.
254    pub padding: (i32, i32),
255
256    /// Dilation. Default to [`Conv2d::DEFAULT_DILATION`] if not specified.
257    pub dilation: (i32, i32),
258
259    /// Groups. Default to [`Conv2d::DEFAULT_GROUPS`] if not specified.
260    pub groups: i32,
261}
262
263impl Conv2d {
264    /// Default value for `with_bias` if not specified.
265    pub const DEFAULT_BIAS: bool = true;
266
267    /// Default value for `stride` if not specified.
268    pub const DEFAULT_STRIDE: SingleOrPair = SingleOrPair::Pair(1, 1);
269
270    /// Default value for `padding` if not specified.
271    pub const DEFAULT_PADDING: SingleOrPair = SingleOrPair::Pair(0, 0);
272
273    /// Default value for `dilation` if not specified.
274    pub const DEFAULT_DILATION: SingleOrPair = SingleOrPair::Pair(1, 1);
275
276    /// Default value for `groups` if not specified.
277    pub const DEFAULT_GROUPS: i32 = 1;
278}
279
280impl Module<&Array> for Conv2d {
281    type Error = Exception;
282    type Output = Array;
283
284    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
285        let mut y = conv2d(
286            x,
287            self.weight.as_ref(),
288            self.stride,
289            self.padding,
290            self.dilation,
291            self.groups,
292        )?;
293        if let Some(bias) = &self.bias.value {
294            y += bias;
295        }
296        Ok(y)
297    }
298
299    fn training_mode(&mut self, _: bool) {}
300}
301
302/// Builder for the `Conv3d` module.
303#[derive(Debug, Clone, Builder)]
304#[builder(
305    root = crate,
306    build_with = build_conv3d,
307    err = ConvBuildError,
308)]
309pub struct Conv3dBuilder {
310    /// Number of input channels.
311    pub input_channels: i32,
312
313    /// Number of output channels.
314    pub output_channels: i32,
315
316    /// Size of the convolution filters.
317    pub kernel_size: SingleOrTriple<i32>,
318
319    /// If `true`, add a learnable bias to the output. Default to [`Conv3d::DEFAULT_BIAS`] if not
320    /// specified.
321    #[builder(optional, default = Conv3d::DEFAULT_BIAS)]
322    pub bias: bool,
323
324    /// Stride. Default to [`Conv3d::DEFAULT_STRIDE`] if not specified.
325    #[builder(optional, default = Conv3d::DEFAULT_STRIDE)]
326    pub stride: SingleOrTriple<i32>,
327
328    /// Padding. Default to [`Conv3d::DEFAULT_PADDING`] if not specified.
329    #[builder(optional, default = Conv3d::DEFAULT_PADDING)]
330    pub padding: SingleOrTriple<i32>,
331
332    /// Dilation. Default to [`Conv3d::DEFAULT_DILATION`] if not specified.
333    #[builder(optional, default = Conv3d::DEFAULT_DILATION)]
334    pub dilation: SingleOrTriple<i32>,
335
336    /// Groups. Default to [`Conv3d::DEFAULT_GROUPS`] if not specified.
337    #[builder(optional, default = Conv3d::DEFAULT_GROUPS)]
338    pub groups: i32,
339}
340
341fn build_conv3d(builder: Conv3dBuilder) -> Result<Conv3d, ConvBuildError> {
342    let input_channels = builder.input_channels;
343    let output_channels = builder.output_channels;
344    let kernel_size: (i32, i32, i32) = builder.kernel_size.into();
345    let with_bias = builder.bias;
346    let padding = builder.padding.into();
347    let stride = builder.stride.into();
348    let dilation = builder.dilation.into();
349    if builder.groups != 1 {
350        return Err(ConvBuildError::GroupedConv3d(builder.groups));
351    }
352
353    let scale =
354        f32::sqrt(1.0 / (input_channels * kernel_size.0 * kernel_size.1 * kernel_size.2) as f32);
355    let weight = uniform::<_, f32>(
356        -scale,
357        scale,
358        &[
359            output_channels,
360            kernel_size.0,
361            kernel_size.1,
362            kernel_size.2,
363            input_channels,
364        ],
365        None,
366    )?;
367    let bias = if with_bias {
368        Some(zeros::<f32>(&[output_channels])?)
369    } else {
370        None
371    };
372
373    Ok(Conv3d {
374        weight: Param::new(weight),
375        bias: Param::new(bias),
376        stride,
377        padding,
378        dilation,
379        groups: builder.groups,
380    })
381}
382
383/// Applies a 3-dimensional convolution over the multi-channel input image.
384///
385/// The channels are expected to be last i.e. the input shape should be `NHWC` where:
386///
387/// - `N` is the batch dimension
388/// - `H` is the input image height
389/// - `W` is the input image width
390/// - `C` is the number of input channels
391#[derive(Debug, Clone, ModuleParameters, Buildable)]
392#[module(root = crate)]
393#[buildable(root = crate)]
394pub struct Conv3d {
395    /// The weight of the convolution layer.
396    #[param]
397    pub weight: Param<Array>,
398
399    /// The bias of the convolution layer.
400    #[param]
401    pub bias: Param<Option<Array>>,
402
403    /// Stride. Default to `(1, 1, 1)` if not specified.
404    pub stride: (i32, i32, i32),
405
406    /// Padding. Default to `(0, 0, 0)` if not specified.
407    pub padding: (i32, i32, i32),
408
409    /// Dilation. Default to `(1, 1, 1)` if not specified.
410    pub dilation: (i32, i32, i32),
411
412    /// Groups. Default to 1 if not specified.
413    pub groups: i32,
414}
415
416impl Conv3d {
417    /// Default value for `with_bias` if not specified.
418    pub const DEFAULT_BIAS: bool = true;
419
420    /// Default value for `stride` if not specified.
421    pub const DEFAULT_STRIDE: SingleOrTriple<i32> = SingleOrTriple::Triple(1, 1, 1);
422
423    /// Default value for `padding` if not specified.
424    pub const DEFAULT_PADDING: SingleOrTriple<i32> = SingleOrTriple::Triple(0, 0, 0);
425
426    /// Default value for `dilation` if not specified.
427    pub const DEFAULT_DILATION: SingleOrTriple<i32> = SingleOrTriple::Triple(1, 1, 1);
428
429    /// Default value for `groups` if not specified.
430    pub const DEFAULT_GROUPS: i32 = 1;
431}
432
433impl Module<&Array> for Conv3d {
434    type Error = Exception;
435    type Output = Array;
436
437    fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
438        let mut y = crate::ops::conv3d(
439            x,
440            self.weight.as_ref(),
441            self.stride,
442            self.padding,
443            self.dilation,
444            self.groups,
445        )?;
446        if let Some(bias) = &self.bias.value {
447            y += bias;
448        }
449        Ok(y)
450    }
451
452    fn training_mode(&mut self, _: bool) {}
453}
454
455/// Validates `groups` for a grouped convolution: it must be positive and divide both channel
456/// counts, matching the weight shape `[C_out, ..., C_in / groups]` that MLX expects.
457fn checked_groups(
458    input_channels: i32,
459    output_channels: i32,
460    groups: i32,
461) -> Result<i32, ConvBuildError> {
462    if groups < 1 || input_channels % groups != 0 || output_channels % groups != 0 {
463        return Err(ConvBuildError::InvalidGroups {
464            groups,
465            input_channels,
466            output_channels,
467        });
468    }
469    Ok(groups)
470}
471
472// The following tests are ported from the swift bindings:
473// mlx-swift/Tests/MLXTests/IntegrationTests.swift
474#[cfg(test)]
475mod tests {
476    use crate::module::Module;
477    use crate::{random::uniform, Dtype};
478    use float_eq::assert_float_eq;
479
480    use crate::nn::Conv1d;
481
482    #[test]
483    fn test_conv1d() {
484        crate::random::seed(819).unwrap();
485        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
486        assert_eq!(a.shape(), &[2, 8, 16]);
487        assert_eq!(a.dtype(), Dtype::Float32);
488        assert_float_eq!(
489            a.mean(None).unwrap().item_exact::<f32>(),
490            0.512_987_5,
491            abs <= 0.010_259_75
492        );
493        assert_float_eq!(
494            a.sum(None).unwrap().item_exact::<f32>(),
495            131.324_8,
496            abs <= 2.626_496
497        );
498        let result = Conv1d::new(16, 2, 8).unwrap().forward(&a).unwrap();
499        assert_eq!(result.shape(), &[2, 1, 2]);
500        assert_eq!(result.dtype(), Dtype::Float32);
501        assert_float_eq!(
502            result.mean(None).unwrap().item_exact::<f32>(),
503            0.264_865_2,
504            abs <= 0.005_297_303_7
505        );
506        assert_float_eq!(
507            result.sum(None).unwrap().item_exact::<f32>(),
508            1.059_460_8,
509            abs <= 0.021_189_215
510        );
511    }
512
513    #[test]
514    fn test_conv2d() {
515        crate::random::seed(62).unwrap();
516        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 8, 4], None).unwrap();
517        assert_eq!(a.shape(), &[2, 8, 8, 4]);
518        assert_eq!(a.dtype(), Dtype::Float32);
519        assert_float_eq!(
520            a.mean(None).unwrap().item_exact::<f32>(),
521            0.522_504_27,
522            abs <= 0.010_450_086
523        );
524        assert_float_eq!(
525            a.sum(None).unwrap().item_exact::<f32>(),
526            267.522_2,
527            abs <= 5.350_444
528        );
529        let result = crate::nn::Conv2d::new(4, 2, (8, 8))
530            .unwrap()
531            .forward(&a)
532            .unwrap();
533        assert_eq!(result.shape(), &[2, 1, 1, 2]);
534        assert_eq!(result.dtype(), Dtype::Float32);
535        assert_float_eq!(
536            result.mean(None).unwrap().item_exact::<f32>(),
537            -0.279_321_5,
538            abs <= 0.005_586_43
539        );
540        assert_float_eq!(
541            result.sum(None).unwrap().item_exact::<f32>(),
542            -1.117_286,
543            abs <= 0.022_345_72
544        );
545    }
546
547    #[test]
548    fn test_conv1d_depthwise_matches_mlx() {
549        use crate::{builder::Builder, nn::Conv1dBuilder, ops::arange};
550
551        let x = &arange::<_, f32>(None, 48, None)
552            .unwrap()
553            .reshape(&[2, 6, 4])
554            .unwrap()
555            / 10.0;
556        let w = &(&arange::<_, f32>(None, 12, None)
557            .unwrap()
558            .reshape(&[4, 3, 1])
559            .unwrap()
560            / 7.0)
561            - 0.5;
562
563        let mut conv = Conv1dBuilder::new(4, 4, 3)
564            .padding(1)
565            .groups(4)
566            .build()
567            .unwrap();
568        assert_eq!(conv.weight.shape(), &[4, 3, 1]);
569        conv.weight.value = w;
570        let y = conv.forward(&x).unwrap();
571        assert_eq!(y.shape(), &[2, 6, 4]);
572
573        // Python MLX 0.32.2: `mx.conv1d(x, w, padding=1, groups=4)`
574        #[rustfmt::skip]
575        let expected: [f32; 48] = [
576            -0.08571428, 0.1142857, 0.4857143, 1.028571,
577            -0.3142857, 0.2214286, 1.014286, 2.064286,
578            -0.7428572, 0.3071429, 1.614286, 3.178571,
579            -1.171429, 0.3928572, 2.214286, 4.292857,
580            -1.6, 0.4785715, 2.814286, 5.407143,
581            -1.514286, 0.02857148, 1.742857, 3.628572,
582            -1.457143, 0.8000001, 3.228571, 5.828571,
583            -2.885714, 0.7357144, 4.614286, 8.75,
584            -3.314286, 0.8214287, 5.214286, 9.864285,
585            -3.742857, 0.907143, 5.814286, 10.97857,
586            -4.171429, 0.9928573, 6.414286, 12.09286,
587            -3.571429, 0.02857158, 3.8, 7.742857,
588        ];
589        for (got, want) in y.to_vec_exact::<f32>().unwrap().iter().zip(expected) {
590            assert_float_eq!(*got, want, abs <= 1e-5 * (1.0 + want.abs()));
591        }
592    }
593
594    #[test]
595    fn test_conv2d_grouped_weight_shape() {
596        use crate::{builder::Builder, nn::Conv2dBuilder};
597
598        let mut conv = Conv2dBuilder::new(4, 8, 3).groups(2).build().unwrap();
599        assert_eq!(conv.weight.shape(), &[8, 3, 3, 2]);
600        let x = uniform::<_, f32>(0.0, 1.0, &[1, 5, 5, 4], None).unwrap();
601        assert_eq!(conv.forward(&x).unwrap().shape(), &[1, 3, 3, 8]);
602    }
603
604    #[test]
605    fn test_conv_groups_must_divide_channels() {
606        use crate::{
607            builder::Builder,
608            error::ConvBuildError,
609            nn::{Conv1dBuilder, Conv2dBuilder, Conv3dBuilder},
610        };
611
612        let invalid = |groups, input_channels, output_channels| ConvBuildError::InvalidGroups {
613            groups,
614            input_channels,
615            output_channels,
616        };
617        assert_eq!(
618            Conv1dBuilder::new(6, 4, 3).groups(4).build().unwrap_err(),
619            invalid(4, 6, 4)
620        );
621        assert_eq!(
622            Conv1dBuilder::new(4, 6, 3).groups(4).build().unwrap_err(),
623            invalid(4, 4, 6)
624        );
625        assert_eq!(
626            Conv1dBuilder::new(4, 4, 3).groups(0).build().unwrap_err(),
627            invalid(0, 4, 4)
628        );
629        assert_eq!(
630            Conv2dBuilder::new(4, 6, 3).groups(4).build().unwrap_err(),
631            invalid(4, 4, 6)
632        );
633        assert_eq!(
634            Conv3dBuilder::new(4, 4, 3).groups(2).build().unwrap_err(),
635            ConvBuildError::GroupedConv3d(2)
636        );
637    }
638}
mlx_rs/nn/convolution.rs

mlx_rs/nn/
convolution.rs
