zsfm_burn/
linear.rs

1use burn::tensor::Tensor;
2
3#[cfg(test)]
4use burn::tensor::{Device, TensorData};
5
6/// y = x @ w^T + b.
7/// `x`: `[lead, d_in]`, `w`: `[d_out, d_in]`, `b`: `[d_out]`.
8pub fn burn_linear_bias(x: Tensor<2>, w: Tensor<2>, b: Tensor<1>) -> Tensor<2> {
9    let y: Tensor<2> = x.matmul(w.transpose());
10    y + b.unsqueeze_dim::<2>(0)
11}
12
13/// y = x @ w^T (no bias).
14pub fn burn_linear_nobias(x: Tensor<2>, w: Tensor<2>) -> Tensor<2> {
15    x.matmul(w.transpose())
16}
17
18/// `y = x @ w^T + b` for rank-D input: flatten leading dims, apply
19/// [`burn_linear_bias`] or [`burn_linear_nobias`], restore shape with the
20/// last dim replaced by `w`'s output dim.
21pub fn burn_linear_nd<const D: usize>(
22    x: Tensor<D>,
23    w: Tensor<2>,
24    b: Option<Tensor<1>>,
25) -> Tensor<D> {
26    let dims = x.dims();
27    let last = dims[D - 1];
28    let lead: usize = dims[..D - 1].iter().product();
29    let out_d = w.dims()[0];
30    let y = match b {
31        Some(bias) => burn_linear_bias(x.reshape([lead, last]), w, bias),
32        None => burn_linear_nobias(x.reshape([lead, last]), w),
33    };
34    let mut out_dims = dims;
35    out_dims[D - 1] = out_d;
36    y.reshape(out_dims)
37}
38
39#[cfg(test)]
40mod tests {
41    use super::*;
42
43    fn dev() -> Device {
44        Device::flex()
45    }
46
47    #[test]
48    fn linear_matches_manual() {
49        let device = dev();
50        let x = Tensor::<2>::from_data(
51            TensorData::from([[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]),
52            &device,
53        );
54        let w = Tensor::<2>::from_data(
55            TensorData::from([[1.0f32, 0.0, 1.0], [0.0, 1.0, 0.0]]),
56            &device,
57        );
58        let b = Tensor::<1>::from_data(TensorData::from([10.0f32, 20.0]), &device);
59        let y = burn_linear_bias(x, w, b);
60        let data: Vec<f32> = y.to_data().try_to_vec().unwrap();
61        assert_eq!(data, vec![14.0, 22.0, 20.0, 25.0]);
62    }
63}
zsfm_burn/linear.rs

zsfm_burn/
linear.rs
