mlx_rs/ops/
convolution.rs

1use crate::error::Result;
2use crate::utils::guard::Guarded;
3use crate::utils::IntoOption;
4use crate::{Array, Stream};
5use mlx_internal_macros::generate_macro;
6
7/// General convolution over an input with several channels returning an error if the inputs are invalid.
8///
9/// - Only 1d and 2d convolutions are supported at the moment
10/// - the default `groups: 1` is currently supported
11///
12/// # Params
13///
14/// - array: Input array of shape `&[N, ..., C_in]`
15/// - weight: Weight array of shape `&[C_out, ..., C_in]`
16/// - strides: The kernel strides. All dimensions get the same stride if only one number is specified.
17/// - padding: The input padding. All dimensions get the same padding if only one number is specified.
18/// - kernel_dilation: The kernel dilation. All dimensions get the same dilation if only one number is specified.
19/// - input_dilation: The input dilation. All dimensions get the same dilation if only one number is specified.
20/// - groups: Input feature groups
21/// - flip: Flip the order in which the spatial dimensions of the weights are processed.
22///   Performs the cross-correlation operator when `flip` is `false` and the convolution
23///   operator otherwise.
24#[allow(clippy::too_many_arguments)]
25pub fn conv_general<'a>(
26    array: impl AsRef<Array>,
27    weight: impl AsRef<Array>,
28    strides: impl IntoOption<&'a [i32]>,
29    padding: impl IntoOption<&'a [i32]>,
30    kernel_dilation: impl IntoOption<&'a [i32]>,
31    input_dilation: impl IntoOption<&'a [i32]>,
32    groups: impl Into<Option<i32>>,
33    flip: impl Into<Option<bool>>,
34) -> Result<Array> {
35    let stream = Stream::thread_local_or_default();
36    let strides = strides.into_option().unwrap_or(&[1]);
37    let padding = padding.into_option().unwrap_or(&[0]);
38    let kernel_dilation = kernel_dilation.into_option().unwrap_or(&[1]);
39    let input_dilation = input_dilation.into_option().unwrap_or(&[1]);
40    let groups = groups.into().unwrap_or(1);
41    let flip = flip.into().unwrap_or(false);
42
43    Array::try_from_op(|res| unsafe {
44        mlx_sys::mlx_conv_general(
45            res,
46            array.as_ref().as_ptr(),
47            weight.as_ref().as_ptr(),
48            strides.as_ptr(),
49            strides.len(),
50            padding.as_ptr(),
51            padding.len(),
52            padding.as_ptr(),
53            padding.len(),
54            kernel_dilation.as_ptr(),
55            kernel_dilation.len(),
56            input_dilation.as_ptr(),
57            input_dilation.len(),
58            groups,
59            flip,
60            stream.as_ref().as_ptr(),
61        )
62    })
63}
64
65/// Compatibility shim for [`conv_general`].
66#[generate_macro(customize(forwarding_shim = true))]
67#[allow(clippy::too_many_arguments)]
68#[deprecated(
69    since = "0.26.0",
70    note = "use `with_stream` or `with_device` around `conv_general`"
71)]
72pub fn conv_general_device<'a>(
73    array: impl AsRef<Array>,
74    weight: impl AsRef<Array>,
75    #[optional] strides: impl IntoOption<&'a [i32]>,
76    #[optional] padding: impl IntoOption<&'a [i32]>,
77    #[optional] kernel_dilation: impl IntoOption<&'a [i32]>,
78    #[optional] input_dilation: impl IntoOption<&'a [i32]>,
79    #[optional] groups: impl Into<Option<i32>>,
80    #[optional] flip: impl Into<Option<bool>>,
81    #[optional] stream: impl AsRef<Stream>,
82) -> Result<Array> {
83    crate::with_stream(stream.as_ref(), || {
84        conv_general(
85            array,
86            weight,
87            strides,
88            padding,
89            kernel_dilation,
90            input_dilation,
91            groups,
92            flip,
93        )
94    })
95}
96
97/// 1D convolution over an input with several channels returning an error if the inputs are invalid.
98///
99/// Only the default `groups=1` is currently supported.
100///
101/// # Params
102///
103/// - array: input array of shape `&[N, H, C_in]`
104/// - weight: weight array of shape `&[C_out, H, C_in]`
105/// - stride: kernel stride. Default to 1 if not specified.
106/// - padding: input padding. Default to 0 if not specified.
107/// - dilation: kernel dilation. Default to 1 if not specified.
108/// - groups: input feature groups. Default to 1 if not specified.
109pub fn conv1d(
110    array: impl AsRef<Array>,
111    weight: impl AsRef<Array>,
112    stride: impl Into<Option<i32>>,
113    padding: impl Into<Option<i32>>,
114    dilation: impl Into<Option<i32>>,
115    groups: impl Into<Option<i32>>,
116) -> Result<Array> {
117    let stream = Stream::thread_local_or_default();
118    let stride = stride.into().unwrap_or(1);
119    let padding = padding.into().unwrap_or(0);
120    let dilation = dilation.into().unwrap_or(1);
121    let groups = groups.into().unwrap_or(1);
122
123    Array::try_from_op(|res| unsafe {
124        mlx_sys::mlx_conv1d(
125            res,
126            array.as_ref().as_ptr(),
127            weight.as_ref().as_ptr(),
128            stride,
129            padding,
130            dilation,
131            groups,
132            stream.as_ref().as_ptr(),
133        )
134    })
135}
136
137/// Compatibility shim for [`conv1d`].
138#[generate_macro(customize(forwarding_shim = true))]
139#[deprecated(
140    since = "0.26.0",
141    note = "use `with_stream` or `with_device` around `conv1d`"
142)]
143pub fn conv1d_device(
144    array: impl AsRef<Array>,
145    weight: impl AsRef<Array>,
146    #[optional] stride: impl Into<Option<i32>>,
147    #[optional] padding: impl Into<Option<i32>>,
148    #[optional] dilation: impl Into<Option<i32>>,
149    #[optional] groups: impl Into<Option<i32>>,
150    #[optional] stream: impl AsRef<Stream>,
151) -> Result<Array> {
152    crate::with_stream(stream.as_ref(), || {
153        conv1d(array, weight, stride, padding, dilation, groups)
154    })
155}
156
157/// 2D convolution over an input with several channels returning an error if the inputs are invalid.
158///
159/// Only the default `groups=1` is currently supported.
160///
161/// # Params
162///
163/// - array: input array of shape `[N, H, W, C_in]`
164/// - weight: weight array of shape `[C_out, H, W, C_in]`
165/// - stride: kernel stride. Default to (1, 1) if not specified.
166/// - padding: input padding. Default to (0, 0) if not specified.
167/// - dilation: kernel dilation. Default to (1, 1) if not specified.
168/// - groups: input feature groups. Default to 1 if not specified.
169pub fn conv2d(
170    array: impl AsRef<Array>,
171    weight: impl AsRef<Array>,
172    stride: impl Into<Option<(i32, i32)>>,
173    padding: impl Into<Option<(i32, i32)>>,
174    dilation: impl Into<Option<(i32, i32)>>,
175    groups: impl Into<Option<i32>>,
176) -> Result<Array> {
177    let stream = Stream::thread_local_or_default();
178    let stride = stride.into().unwrap_or((1, 1));
179    let padding = padding.into().unwrap_or((0, 0));
180    let dilation = dilation.into().unwrap_or((1, 1));
181    let groups = groups.into().unwrap_or(1);
182
183    Array::try_from_op(|res| unsafe {
184        mlx_sys::mlx_conv2d(
185            res,
186            array.as_ref().as_ptr(),
187            weight.as_ref().as_ptr(),
188            stride.0,
189            stride.1,
190            padding.0,
191            padding.1,
192            dilation.0,
193            dilation.1,
194            groups,
195            stream.as_ref().as_ptr(),
196        )
197    })
198}
199
200/// Compatibility shim for [`conv2d`].
201#[generate_macro(customize(forwarding_shim = true))]
202#[deprecated(
203    since = "0.26.0",
204    note = "use `with_stream` or `with_device` around `conv2d`"
205)]
206pub fn conv2d_device(
207    array: impl AsRef<Array>,
208    weight: impl AsRef<Array>,
209    #[optional] stride: impl Into<Option<(i32, i32)>>,
210    #[optional] padding: impl Into<Option<(i32, i32)>>,
211    #[optional] dilation: impl Into<Option<(i32, i32)>>,
212    #[optional] groups: impl Into<Option<i32>>,
213    #[optional] stream: impl AsRef<Stream>,
214) -> Result<Array> {
215    crate::with_stream(stream.as_ref(), || {
216        conv2d(array, weight, stride, padding, dilation, groups)
217    })
218}
219
220/// 3D convolution over an input with several channels.
221///
222/// Only the default `groups=1` is currently supported.
223pub fn conv3d(
224    array: impl AsRef<Array>,
225    weight: impl AsRef<Array>,
226    stride: impl Into<Option<(i32, i32, i32)>>,
227    padding: impl Into<Option<(i32, i32, i32)>>,
228    dilation: impl Into<Option<(i32, i32, i32)>>,
229    groups: impl Into<Option<i32>>,
230) -> Result<Array> {
231    let stream = Stream::thread_local_or_default();
232    let stride = stride.into().unwrap_or((1, 1, 1));
233    let padding = padding.into().unwrap_or((0, 0, 0));
234    let dilation = dilation.into().unwrap_or((1, 1, 1));
235    let groups = groups.into().unwrap_or(1);
236
237    Array::try_from_op(|res| unsafe {
238        mlx_sys::mlx_conv3d(
239            res,
240            array.as_ref().as_ptr(),
241            weight.as_ref().as_ptr(),
242            stride.0,
243            stride.1,
244            stride.2,
245            padding.0,
246            padding.1,
247            padding.2,
248            dilation.0,
249            dilation.1,
250            dilation.2,
251            groups,
252            stream.as_ref().as_ptr(),
253        )
254    })
255}
256
257/// Compatibility shim for [`conv3d`].
258#[generate_macro(customize(forwarding_shim = true))]
259#[deprecated(
260    since = "0.26.0",
261    note = "use `with_stream` or `with_device` around `conv3d`"
262)]
263pub fn conv3d_device(
264    array: impl AsRef<Array>,
265    weight: impl AsRef<Array>,
266    #[optional] stride: impl Into<Option<(i32, i32, i32)>>,
267    #[optional] padding: impl Into<Option<(i32, i32, i32)>>,
268    #[optional] dilation: impl Into<Option<(i32, i32, i32)>>,
269    #[optional] groups: impl Into<Option<i32>>,
270    #[optional] stream: impl AsRef<Stream>,
271) -> Result<Array> {
272    crate::with_stream(stream.as_ref(), || {
273        conv3d(array, weight, stride, padding, dilation, groups)
274    })
275}
276
277/// 1D transposed convolution over an input with several channels.
278///
279/// Only the default `groups=1` is currently supported.
280///
281/// # Params
282///
283/// - array: input array of shape `[N, H, C_in]`
284/// - weight: weight array of shape `[C_out, H, C_in]`
285/// - stride: kernel stride. Default to 1 if not specified.
286/// - padding: input padding. Default to 0 if not specified.
287/// - dilation: kernel dilation. Default to 1 if not specified.
288/// - groups: input feature groups. Default to 1 if not specified.
289/// - stream: stream or device to evaluate on.
290#[allow(clippy::too_many_arguments)]
291pub fn conv_transpose1d(
292    array: impl AsRef<Array>,
293    weight: impl AsRef<Array>,
294    stride: impl Into<Option<i32>>,
295    padding: impl Into<Option<i32>>,
296    dilation: impl Into<Option<i32>>,
297    output_padding: impl Into<Option<i32>>,
298    groups: impl Into<Option<i32>>,
299) -> Result<Array> {
300    let stream = Stream::thread_local_or_default();
301    let stride = stride.into().unwrap_or(1);
302    let padding = padding.into().unwrap_or(0);
303    let dilation = dilation.into().unwrap_or(1);
304    let output_padding = output_padding.into().unwrap_or(0);
305    let groups = groups.into().unwrap_or(1);
306
307    Array::try_from_op(|res| unsafe {
308        mlx_sys::mlx_conv_transpose1d(
309            res,
310            array.as_ref().as_ptr(),
311            weight.as_ref().as_ptr(),
312            stride,
313            padding,
314            dilation,
315            output_padding,
316            groups,
317            stream.as_ref().as_ptr(),
318        )
319    })
320}
321
322/// Compatibility shim for [`conv_transpose1d`].
323#[allow(clippy::too_many_arguments)]
324#[generate_macro(customize(forwarding_shim = true))]
325#[deprecated(
326    since = "0.26.0",
327    note = "use `with_stream` or `with_device` around `conv_transpose1d`"
328)]
329pub fn conv_transpose1d_device(
330    array: impl AsRef<Array>,
331    weight: impl AsRef<Array>,
332    #[optional] stride: impl Into<Option<i32>>,
333    #[optional] padding: impl Into<Option<i32>>,
334    #[optional] dilation: impl Into<Option<i32>>,
335    #[optional] output_padding: impl Into<Option<i32>>,
336    #[optional] groups: impl Into<Option<i32>>,
337    #[optional] stream: impl AsRef<Stream>,
338) -> Result<Array> {
339    crate::with_stream(stream.as_ref(), || {
340        conv_transpose1d(
341            array,
342            weight,
343            stride,
344            padding,
345            dilation,
346            output_padding,
347            groups,
348        )
349    })
350}
351
352/// 2D transposed convolution over an input with several channels.
353///
354/// Only the default `groups=1` is currently supported.
355///
356/// The numeric parameters may be given as single values:
357///
358/// # Params
359/// - array: input array of shape `[N, H, W, C_in]`
360/// - weight: weight array of shape `[C_out, H, W, C_in]`
361/// - stride: kernel stride. Default to (1, 1) if not specified.
362/// - padding: input padding. Default to (0, 0) if not specified.
363/// - dilation: kernel dilation. Default to (1, 1) if not specified.
364/// - groups: input feature groups. Default to 1 if not specified.
365/// - stream: stream or device to evaluate on.
366#[allow(clippy::too_many_arguments)]
367pub fn conv_transpose2d(
368    array: impl AsRef<Array>,
369    weight: impl AsRef<Array>,
370    stride: impl Into<Option<(i32, i32)>>,
371    padding: impl Into<Option<(i32, i32)>>,
372    dilation: impl Into<Option<(i32, i32)>>,
373    output_padding: impl Into<Option<(i32, i32)>>,
374    groups: impl Into<Option<i32>>,
375) -> Result<Array> {
376    let stream = Stream::thread_local_or_default();
377    let stride = stride.into().unwrap_or((1, 1));
378    let padding = padding.into().unwrap_or((0, 0));
379    let dilation = dilation.into().unwrap_or((1, 1));
380    let output_padding = output_padding.into().unwrap_or((0, 0));
381    let groups = groups.into().unwrap_or(1);
382
383    Array::try_from_op(|res| unsafe {
384        mlx_sys::mlx_conv_transpose2d(
385            res,
386            array.as_ref().as_ptr(),
387            weight.as_ref().as_ptr(),
388            stride.0,
389            stride.1,
390            padding.0,
391            padding.1,
392            dilation.0,
393            dilation.1,
394            output_padding.0,
395            output_padding.1,
396            groups,
397            stream.as_ref().as_ptr(),
398        )
399    })
400}
401
402/// Compatibility shim for [`conv_transpose2d`].
403#[allow(clippy::too_many_arguments)]
404#[generate_macro(customize(forwarding_shim = true))]
405#[deprecated(
406    since = "0.26.0",
407    note = "use `with_stream` or `with_device` around `conv_transpose2d`"
408)]
409pub fn conv_transpose2d_device(
410    array: impl AsRef<Array>,
411    weight: impl AsRef<Array>,
412    #[optional] stride: impl Into<Option<(i32, i32)>>,
413    #[optional] padding: impl Into<Option<(i32, i32)>>,
414    #[optional] dilation: impl Into<Option<(i32, i32)>>,
415    #[optional] output_padding: impl Into<Option<(i32, i32)>>,
416    #[optional] groups: impl Into<Option<i32>>,
417    #[optional] stream: impl AsRef<Stream>,
418) -> Result<Array> {
419    crate::with_stream(stream.as_ref(), || {
420        conv_transpose2d(
421            array,
422            weight,
423            stride,
424            padding,
425            dilation,
426            output_padding,
427            groups,
428        )
429    })
430}
431
432/// 3D transposed convolution over an input with several channels.
433///
434/// Only the default `groups=1` is currently supported.
435///
436/// The numeric parameters may be given as single values:
437///
438/// # Params
439/// - array: input array of shape `[N, D, H, W, C_in]`
440/// - weight: weight array of shape `[C_out, D, H, W, C_in]`
441/// - stride: kernel stride. Default to (1, 1, 1) if not specified.
442/// - padding: input padding. Default to (0, 0, 0) if not specified.
443/// - dilation: kernel dilation. Default to (1, 1, 1) if not specified.
444/// - groups: input feature groups. Default to 1 if not specified.
445/// - stream: stream or device to evaluate on.
446#[allow(clippy::too_many_arguments)]
447pub fn conv_transpose3d(
448    array: impl AsRef<Array>,
449    weight: impl AsRef<Array>,
450    stride: impl Into<Option<(i32, i32, i32)>>,
451    padding: impl Into<Option<(i32, i32, i32)>>,
452    dilation: impl Into<Option<(i32, i32, i32)>>,
453    output_padding: impl Into<Option<(i32, i32, i32)>>,
454    groups: impl Into<Option<i32>>,
455) -> Result<Array> {
456    let stream = Stream::thread_local_or_default();
457    let stride = stride.into().unwrap_or((1, 1, 1));
458    let padding = padding.into().unwrap_or((0, 0, 0));
459    let dilation = dilation.into().unwrap_or((1, 1, 1));
460    let output_padding = output_padding.into().unwrap_or((0, 0, 0));
461    let groups = groups.into().unwrap_or(1);
462
463    Array::try_from_op(|res| unsafe {
464        mlx_sys::mlx_conv_transpose3d(
465            res,
466            array.as_ref().as_ptr(),
467            weight.as_ref().as_ptr(),
468            stride.0,
469            stride.1,
470            stride.2,
471            padding.0,
472            padding.1,
473            padding.2,
474            dilation.0,
475            dilation.1,
476            dilation.2,
477            output_padding.0,
478            output_padding.1,
479            output_padding.2,
480            groups,
481            stream.as_ref().as_ptr(),
482        )
483    })
484}
485
486/// Compatibility shim for [`conv_transpose3d`].
487#[allow(clippy::too_many_arguments)]
488#[generate_macro(customize(forwarding_shim = true))]
489#[deprecated(
490    since = "0.26.0",
491    note = "use `with_stream` or `with_device` around `conv_transpose3d`"
492)]
493pub fn conv_transpose3d_device(
494    array: impl AsRef<Array>,
495    weight: impl AsRef<Array>,
496    #[optional] stride: impl Into<Option<(i32, i32, i32)>>,
497    #[optional] padding: impl Into<Option<(i32, i32, i32)>>,
498    #[optional] dilation: impl Into<Option<(i32, i32, i32)>>,
499    #[optional] output_padding: impl Into<Option<(i32, i32, i32)>>,
500    #[optional] groups: impl Into<Option<i32>>,
501    #[optional] stream: impl AsRef<Stream>,
502) -> Result<Array> {
503    crate::with_stream(stream.as_ref(), || {
504        conv_transpose3d(
505            array,
506            weight,
507            stride,
508            padding,
509            dilation,
510            output_padding,
511            groups,
512        )
513    })
514}
515
516#[cfg(test)]
517mod tests {
518    use super::*;
519    use crate::test_utils::{assert_array_eq, tolerances};
520    use pretty_assertions::assert_eq;
521
522    #[test]
523    fn test_conv1d_complex_device() {
524        // Define a 1D input with two channels
525        let input_data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
526        let input_array = Array::from_slice(&input_data, &[1, 5, 2]);
527
528        // Define a 1D kernel with two input channels and two output channels
529        let weight_data = [0.5, 0.0, -0.5, 1.0, 0.0, 1.5, 2.0, 0.0, -2.0, 1.5, 0.0, 1.0];
530        let weight_array = Array::from_slice(&weight_data, &[2, 3, 2]);
531
532        let result = conv1d(
533            &input_array,
534            &weight_array,
535            Some(1), // stride
536            Some(0), // padding
537            Some(1), // dilation
538            Some(1), // groups
539        )
540        .unwrap();
541
542        let expected_output = [12.0, 8.0, 17.0, 13.0, 22.0, 18.0];
543        assert_eq!(result.shape(), &[1, 3, 2]);
544        assert_eq!(result.as_slice::<f32>(), &expected_output);
545    }
546
547    #[test]
548    fn test_conv_transpose1d() {
549        // Single channel input
550        let input = Array::from_slice(&[1.0, 2.0, 3.0], &[1, 3, 1]);
551        // Single input/output channel kernel
552        let weights = Array::from_slice(&[1.0, 0.5], &[1, 2, 1]);
553
554        let result = conv_transpose1d(
555            &input,
556            &weights,
557            Some(1), // stride
558            Some(0), // padding
559            Some(1), // dilation
560            None,    // output padding
561            Some(1), // groups
562        )
563        .unwrap();
564
565        let expected = [1.0, 2.5, 4.0, 1.5];
566        assert_eq!(result.shape(), &[1, 4, 1]);
567        assert_eq!(result.as_slice::<f32>(), &expected);
568    }
569
570    #[test]
571    fn test_conv2d() {
572        // Define a 2x2 input with one channel (grayscale image or similar)
573        let input_data = [1.0, 2.0, 3.0, 4.0];
574        let input_shape = [1, 2, 2, 1]; // [N, H, W, C]
575        let input_array = Array::from_slice(&input_data, &input_shape);
576
577        // Define a 2x2 kernel with one input channel and one output channel
578        let weight_data = [1.0, 0.0, 0.0, 1.0];
579        let weight_shape = [1, 2, 2, 1]; // [C_out, H_k, W_k, C_in]
580        let weight_array = Array::from_slice(&weight_data, &weight_shape);
581
582        // Perform the convolution with no padding and stride of 1
583        let result = conv2d(
584            &input_array,
585            &weight_array,
586            Some((1, 1)), // stride
587            Some((0, 0)), // padding
588            Some((1, 1)), // dilation
589            Some(1),      // groups
590        )
591        .unwrap();
592
593        // Expected result is the convolution of a 2x2 filter over a 2x2 input with valid padding, resulting in a single output value
594        let expected_output = 1.0 * 1.0 + 2.0 * 0.0 + 3.0 * 0.0 + 4.0 * 1.0; // = 1*1 + 4*1 = 5
595        assert_eq!(result.as_slice::<f32>(), &[expected_output]);
596    }
597
598    #[test]
599    fn test_conv_transpose2d() {
600        // 2x2 single channel input
601        let input = Array::from_slice(&[1.0, 2.0, 3.0, 4.0], &[1, 2, 2, 1]);
602        // 2x2 single channel kernel (identity-like)
603        let weights = Array::from_slice(&[1.0, 0.0, 0.0, 1.0], &[1, 2, 2, 1]);
604
605        let result = conv_transpose2d(
606            &input,
607            &weights,
608            Some((1, 1)), // stride
609            Some((0, 0)), // padding
610            Some((1, 1)), // dilation
611            None,         // output padding
612            Some(1),      // groups
613        )
614        .unwrap();
615
616        let expected = [1.0, 2.0, 0.0, 3.0, 5.0, 2.0, 0.0, 3.0, 4.0];
617        assert_eq!(result.shape(), &[1, 3, 3, 1]);
618        assert_eq!(result.as_slice::<f32>(), &expected);
619    }
620
621    #[test]
622    fn test_conv3d() {
623        // Define a 2x2x2 input with one channel
624        let input_data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
625        let input_shape = [1, 2, 2, 2, 1]; // [N, D, H, W, C]
626        let input_array = Array::from_slice(&input_data, &input_shape);
627
628        // Define a 2x2x2 kernel with one input channel and one output channel
629        let weight_data = [1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
630        let weight_shape = [1, 2, 2, 2, 1]; // [C_out, D_k, H_k, W_k, C_in]
631        let weight_array = Array::from_slice(&weight_data, &weight_shape);
632
633        // Perform the convolution with no padding and stride of 1
634        let result = conv3d(
635            &input_array,
636            &weight_array,
637            Some((1, 1, 1)), // stride
638            Some((0, 0, 0)), // padding
639            Some((1, 1, 1)), // dilation
640            Some(1),         // groups
641        )
642        .unwrap();
643
644        // Expected result is the convolution of a 2x2x2 filter over a 2x2x2 input with valid padding, resulting in a single output value
645        let expected_output = 1.0 * 1.0
646            + 2.0 * 0.0
647            + 3.0 * 0.0
648            + 4.0 * 1.0
649            + 5.0 * 0.0
650            + 6.0 * 1.0
651            + 7.0 * 1.0
652            + 8.0 * 0.0; // = 1*1 + 4*1 + 6*1 + 7*1 = 18
653        assert_array_eq(
654            result,
655            Array::from_slice(&[expected_output], &[1, 1, 1, 1, 1]),
656            tolerances::EXACT.rtol,
657            tolerances::EXACT.atol,
658        );
659    }
660
661    #[test]
662    fn test_conv_transpose3d() {
663        // 2x2x2 single channel input
664        let input = Array::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[1, 2, 2, 2, 1]);
665        // 2x2x2 single channel kernel
666        let weights =
667            Array::from_slice(&[1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0], &[1, 2, 2, 2, 1]);
668
669        let result = conv_transpose3d(
670            &input,
671            &weights,
672            Some((1, 1, 1)), // stride
673            Some((0, 0, 0)), // padding
674            Some((1, 1, 1)), // dilation
675            None,            // output padding
676            Some(1),         // groups
677        )
678        .unwrap();
679
680        assert_eq!(result.shape(), &[1, 3, 3, 3, 1]);
681    }
682
683    #[test]
684    fn test_conv_wrong_dimensions() {
685        let input_data = [1.0, 2.0, 3.0, 4.0];
686        let input_shape = [1, 2, 2, 1]; // [N, H, W, C]
687        let input_array = Array::from_slice(&input_data, &input_shape);
688
689        let weight_data = [1.0, 0.0, 0.0, 1.0];
690        let weight_shape = [1, 2, 2]; // [C_out, H_k, W_k]
691        let weight_array = Array::from_slice(&weight_data, &weight_shape);
692
693        let result = conv2d(
694            &input_array,
695            &weight_array,
696            Some((1, 1)), // stride
697            Some((0, 0)), // padding
698            Some((1, 1)), // dilation
699            Some(1),      // groups
700        );
701
702        assert!(result.is_err());
703    }
704
705    #[test]
706    fn test_conv_invalid_group_size() {
707        let input_data = [1.0, 2.0, 3.0, 4.0];
708        let input_shape = [1, 2, 2, 1]; // [N, H, W, C]
709        let input_array = Array::from_slice(&input_data, &input_shape);
710
711        let weight_data = [1.0, 0.0, 0.0, 1.0];
712        let weight_shape = [1, 2, 2, 1]; // [C_out, H_k, W_k, C_in]
713        let weight_array = Array::from_slice(&weight_data, &weight_shape);
714
715        let result = conv2d(
716            &input_array,
717            &weight_array,
718            Some((1, 1)), // stride
719            Some((0, 0)), // padding
720            Some((1, 1)), // dilation
721            Some(2),      // groups
722        );
723
724        assert!(result.is_err());
725    }
726
727    #[test]
728    fn test_conv_non_float() {
729        let input_data = [1, 2, 3, 4];
730        let input_shape = [1, 2, 2, 1]; // [N, H, W, C]
731        let input_array = Array::from_slice(&input_data, &input_shape);
732
733        let weight_data = [1, 0, 0, 1];
734        let weight_shape = [1, 2, 2, 1]; // [C_out, H_k, W_k, C_in]
735        let weight_array = Array::from_slice(&weight_data, &weight_shape);
736
737        let result = conv2d(
738            &input_array,
739            &weight_array,
740            Some((1, 1)), // stride
741            Some((0, 0)), // padding
742            Some((1, 1)), // dilation
743            Some(1),      // groups
744        );
745
746        assert!(result.is_err());
747    }
748}
mlx_rs/ops/convolution.rs

mlx_rs/ops/
convolution.rs
