nuvec_core/
conv.rs

1//! Convolutions, as MLX's `conv1d`.
2
3use crate::ops::{add, matmul, multiply, reshape, slice, transpose};
4use crate::{
5    Array, Category, DType, Error, Result, SliceArg, concatenate, promote_types, stack, zeros,
6};
7
8/// The 1-D convolution (cross-correlation) of `input` (`[N, L, C_in]`) with `weight`
9/// (`[C_out, K, C_in / groups]`), as MLX's `conv1d`: output `[N, L_out, C_out]` with
10/// `L_out = (L + 2·padding − dilation·(K − 1) − 1) / stride + 1`. Channels split into
11/// `groups` independent groups (`groups = C_in = C_out` is a depthwise convolution).
12/// Products accumulate in float32 and round once, as MLX's kernels.
13///
14/// ```
15/// use nuvec_core::{Array, Result, conv1d};
16///
17/// # fn main() -> Result<()> {
18/// // A moving sum of three, one channel.
19/// let x = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0], &[1, 5, 1])?;
20/// let w = Array::from_slice(&[1.0f32, 1.0, 1.0], &[1, 3, 1])?;
21/// assert_eq!(conv1d(&x, &w, 1, 0, 1, 1)?.to_vec::<f32>()?, [6.0, 9.0, 12.0]);
22/// assert_eq!(conv1d(&x, &w, 2, 1, 1, 1)?.to_vec::<f32>()?, [3.0, 9.0, 9.0]);
23/// # Ok(())
24/// # }
25/// ```
26pub fn conv1d(
27    input: &Array,
28    weight: &Array,
29    stride: usize,
30    padding: usize,
31    dilation: usize,
32    groups: usize,
33) -> Result<Array> {
34    let invalid = |msg: String| Err(Error::InvalidArgument(format!("conv1d: {msg}")));
35    if input.ndim() != 3 || weight.ndim() != 3 {
36        return invalid(format!(
37            "input must be [N, L, C_in] and weight [C_out, K, C_in / groups], got {:?} and {:?}",
38            input.shape(),
39            weight.shape()
40        ));
41    }
42    let (n, l, c_in) = (input.shape()[0], input.shape()[1], input.shape()[2]);
43    let (c_out, k, c_group) = (weight.shape()[0], weight.shape()[1], weight.shape()[2]);
44    if stride == 0 || dilation == 0 || groups == 0 {
45        return invalid("stride, dilation and groups must be positive".into());
46    }
47    if c_in % groups != 0 || c_out % groups != 0 || c_group * groups != c_in {
48        return invalid(format!(
49            "{c_in} input and {c_out} output channels in {groups} groups need weights \
50             [C_out, K, {}], got {:?}",
51            c_in / groups.max(1),
52            weight.shape()
53        ));
54    }
55    let span = dilation * (k.max(1) - 1) + 1;
56    if k == 0 || l + 2 * padding < span {
57        return invalid(format!(
58            "a kernel spanning {span} positions does not fit {l} positions padded by {padding}"
59        ));
60    }
61    let dtype = promote_types(input.dtype(), weight.dtype());
62    if dtype.category() != Category::Float {
63        return Err(Error::UnsupportedDType {
64            op: "conv1d",
65            dtype,
66        });
67    }
68    let acc = if dtype == DType::F64 {
69        DType::F64
70    } else {
71        DType::F32
72    };
73    let l_out = (l + 2 * padding - span) / stride + 1;
74    let mut x = input.astype(acc)?;
75    if padding > 0 {
76        let pad = zeros(&[n, padding, c_in], acc)?;
77        x = concatenate(&[pad.clone(), x, pad], 1)?;
78    }
79    // Tap `j` of every output position: positions `j·dilation + i·stride`.
80    let tap = |j: usize| -> Result<Array> {
81        let start = (j * dilation) as isize;
82        let stop = start + ((l_out - 1) * stride + 1) as isize;
83        slice(
84            &x,
85            &[
86                SliceArg::full(),
87                SliceArg::new(Some(start), Some(stop), stride as isize),
88                SliceArg::full(),
89            ],
90        )
91    };
92    let w = weight.astype(acc)?;
93    let out = if groups == c_in && c_out == c_in {
94        // Depthwise: each channel's taps weighted by its own kernel.
95        let mut sum: Option<Array> = None;
96        for j in 0..k {
97            let wj = reshape(
98                &slice(
99                    &w,
100                    &[
101                        SliceArg::full(),
102                        SliceArg::new(Some(j as isize), Some(j as isize + 1), 1),
103                        SliceArg::full(),
104                    ],
105                )?,
106                &[c_in as isize],
107            )?;
108            let term = multiply(&tap(j)?, &wj)?;
109            sum = Some(match sum {
110                Some(s) => add(&s, &term)?,
111                None => term,
112            });
113        }
114        sum.expect("at least one tap")
115    } else {
116        // `[N·L_out, K·C_in]` columns, one product per group.
117        let cols = stack(&(0..k).map(tap).collect::<Result<Vec<_>>>()?, 2)?;
118        let per_out = c_out / groups;
119        let mut outs = Vec::with_capacity(groups);
120        for g in 0..groups {
121            let cg = slice(
122                &cols,
123                &[
124                    SliceArg::full(),
125                    SliceArg::full(),
126                    SliceArg::full(),
127                    SliceArg::new(
128                        Some((g * c_group) as isize),
129                        Some(((g + 1) * c_group) as isize),
130                        1,
131                    ),
132                ],
133            )?;
134            let cg = reshape(&cg, &[(n * l_out) as isize, (k * c_group) as isize])?;
135            let wg = slice(
136                &w,
137                &[
138                    SliceArg::new(
139                        Some((g * per_out) as isize),
140                        Some(((g + 1) * per_out) as isize),
141                        1,
142                    ),
143                    SliceArg::full(),
144                    SliceArg::full(),
145                ],
146            )?;
147            let wg = reshape(&wg, &[per_out as isize, (k * c_group) as isize])?;
148            outs.push(matmul(&cg, &transpose(&wg, &[1, 0])?)?);
149        }
150        let y = match groups {
151            1 => outs.pop().expect("one group"),
152            _ => concatenate(&outs, 1)?,
153        };
154        reshape(&y, &[n as isize, l_out as isize, c_out as isize])?
155    };
156    out.astype(dtype)
157}
nuvec_core/conv.rs

nuvec_core/
conv.rs
