zsfm_nn/
linear.rs

1use anyhow::{Context, Result};
2use candle_core::Tensor;
3
4/// `y = x @ w^T + b`, flattening any leading dims of `x` into a batch dim so this works for
5/// rank-2 or higher inputs against a rank-2 weight `[d_out, d_in]`.
6pub fn linear(x: &Tensor, w: &Tensor, b: Option<&Tensor>) -> Result<Tensor> {
7    let dims = x.dims().to_vec();
8    let d_in = *dims.last().context("linear: input has no dims")?;
9    let lead: usize = dims[..dims.len() - 1].iter().product();
10    let x2 = x.reshape((lead, d_in))?;
11    let y2 = x2.matmul(&w.t()?)?;
12    let d_out = w.dim(0)?;
13    let mut out_dims = dims[..dims.len() - 1].to_vec();
14    out_dims.push(d_out);
15    let y = y2.reshape(out_dims)?;
16    match b {
17        Some(b) => Ok(y.broadcast_add(b)?),
18        None => Ok(y),
19    }
20}
21
22pub fn linear_nobias(x: &Tensor, w: &Tensor) -> Result<Tensor> {
23    linear(x, w, None)
24}
25
26pub fn linear_bias(x: &Tensor, w: &Tensor, b: &Tensor) -> Result<Tensor> {
27    linear(x, w, Some(b))
28}
29
30#[cfg(test)]
31mod tests {
32    use super::*;
33    use candle_core::Device;
34
35    #[test]
36    fn matches_manual_matmul_for_3d_input() {
37        let device = Device::Cpu;
38        let x = Tensor::from_vec((0..24u32).map(|v| v as f32).collect(), (2, 3, 4), &device).unwrap();
39        let w = Tensor::from_vec((0..8u32).map(|v| v as f32 * 0.5).collect(), (2, 4), &device).unwrap();
40        let b = Tensor::from_vec(vec![1.0f32, -1.0], (2,), &device).unwrap();
41
42        let got = linear(&x, &w, Some(&b)).unwrap();
43        assert_eq!(got.dims(), &[2, 3, 2]);
44
45        // Manual reference: flatten to 2D, matmul, add bias, reshape.
46        let x2 = x.reshape((6, 4)).unwrap();
47        let want2 = x2.matmul(&w.t().unwrap()).unwrap().broadcast_add(&b).unwrap();
48        let want = want2.reshape((2, 3, 2)).unwrap();
49
50        let got_v: Vec<f32> = got.flatten_all().unwrap().to_vec1().unwrap();
51        let want_v: Vec<f32> = want.flatten_all().unwrap().to_vec1().unwrap();
52        assert_eq!(got_v, want_v);
53    }
54
55    #[test]
56    fn nobias_matches_direct_matmul_for_2d_input() {
57        let device = Device::Cpu;
58        let x = Tensor::from_vec((0..12u32).map(|v| v as f32).collect(), (3, 4), &device).unwrap();
59        let w = Tensor::from_vec((0..8u32).map(|v| v as f32).collect(), (2, 4), &device).unwrap();
60
61        let got = linear_nobias(&x, &w).unwrap();
62        let want = x.matmul(&w.t().unwrap()).unwrap();
63
64        let got_v: Vec<f32> = got.flatten_all().unwrap().to_vec1().unwrap();
65        let want_v: Vec<f32> = want.flatten_all().unwrap().to_vec1().unwrap();
66        assert_eq!(got_v, want_v);
67    }
68}
zsfm_nn/linear.rs

zsfm_nn/
linear.rs
