Struct Linear 

pub struct Linear {
    pub weight: Array,
    pub bias: Option<Array>,
}

An affine layer y = x Wᵀ + b, with W of shape [output, input] and b of shape [output] (MLX’s nn.Linear).

Fields

weight: Array

The weight, [output_dims, input_dims].

bias: Option<Array>

The bias, [output_dims].

Implementations

impl Linear

pub fn new(input_dims: usize, output_dims: usize, bias: bool) -> Result<Self>

A layer initialised as MLX’s nn.Linear: weight, then bias, drawn uniformly from [-1/√input, 1/√input).

Examples found in repository

examples/mnist.rs (line 33)

    fn new(sizes: &[usize]) -> Result<Self> {
        let layers = sizes
            .windows(2)
            .map(|w| Linear::new(w[0], w[1], true))
            .collect::<Result<_>>()?;
        Ok(Self { layers })
    }

pub fn forward(&self, x: &Array) -> Result<Array>

x Wᵀ + b over the last axis of x; with a bias it is one fused addmm, as in MLX.

Examples found in repository

examples/mnist.rs (line 42)

    fn forward(&self, x: &Array) -> Result<Array> {
        let (last, hidden) = self.layers.split_last().expect("at least one layer");
        let mut h = x.clone();
        for layer in hidden {
            h = relu(&layer.forward(&h)?)?;
        }
        last.forward(&h)
    }

Trait Implementations

impl Clone for Linear

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Linear

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Module for Linear

fn visit(&mut self, p: &mut Parameters<'_>)

Reports every parameter and submodule to p, frozen ones included.

fn parameters(&self) -> Vec<(String, Array)>

Every parameter with its path (layers.0.weight), in visiting order.

fn trainable_parameters(&self) -> Vec<(String, Array)>

The parameters that are trained (not frozen), with their paths.

fn update_trainable(&mut self, values: &[Array]) -> Result<()>

Replaces the trainable parameters, given in the order of Module::trainable_parameters (as optimizers return them).

fn load_weights( &mut self, weights: impl IntoIterator<Item = (String, Array)>, strict: bool, ) -> Result<()>

Replaces parameters by path, as MLX’s load_weights: each weight must have its parameter’s shape (the dtype is kept as loaded). With strict, every parameter must be given and no unknown names are allowed.