zyx_nn/
linear.rs

1// Copyright (C) 2025 zk4x
2// SPDX-License-Identifier: LGPL-3.0-only
3
4use zyx::{DType, Tensor, ZyxError};
5use zyx_derive::Module;
6
7/// Linear layer
8#[derive(Debug, Module)]
9#[cfg_attr(feature = "py", pyo3::pyclass(get_all, set_all))]
10pub struct Linear {
11    /// weight
12    pub weight: Tensor,
13    /// bias
14    pub bias: Option<Tensor>,
15}
16
17impl Linear {
18    /// Initilize linear layer in device self
19    pub fn new(
20        in_features: u64,
21        out_features: u64,
22        bias: bool,
23        dtype: DType,
24    ) -> Result<Linear, ZyxError> {
25        let l = -(1.0 / (in_features as f32)).sqrt();
26        let u = (1.0 / (in_features as f32)).sqrt();
27        Ok(Linear {
28            weight: Tensor::uniform([out_features, in_features], l..u)?.cast(dtype),
29            bias: if bias {
30                Some(Tensor::uniform([out_features], l..u)?.cast(dtype))
31            } else {
32                None
33            },
34        })
35    }
36
37    /// Forward function for linear.
38    /// Calculates x.dot(&self.weight) + self.bias
39    pub fn forward(&self, x: impl Into<Tensor>) -> Result<Tensor, ZyxError> {
40        let x = x.into().dot(self.weight.t())?;
41        if let Some(bias) = &self.bias {
42            return Ok(x + bias);
43        }
44        Ok(x)
45    }
46}
47
48#[test]
49fn linear() -> Result<(), ZyxError> {
50    let l0 = Linear::new(4, 16, true, DType::F32)?;
51    println!("{}\n{}", l0.weight, l0.bias.as_ref().unwrap());
52    let x = Tensor::randn([8, 4], DType::F32)?;
53    let y = l0.forward(x)?.relu();
54
55    println!("{y}");
56
57    Ok(())
58}
zyx_nn/linear.rs

zyx_nn/
linear.rs
