Models

AodhML provides a Model trait and associated types for defining neural network architectures in a type-safe manner.

Defining a Model

import "std:nn"
import "std:tensor"

type Linear = struct {
    weights: Tensor<f32>,
    bias: Tensor<f32>,
}

fn Linear::new(in_features: u32, out_features: u32) -> Linear {
    let scale = sqrt(2.0 / in_features as f32)
    return Linear {
        weights: Tensor::randn([in_features, out_features]) * scale,
        bias: Tensor::zeros([out_features]),
    }
}

fn Linear::forward(self, x: Tensor<f32>) -> Tensor<f32> {
    return matmul(x, self.weights) + self.bias
}

Composing Models

type MLP = struct {
    fc1: Linear,
    fc2: Linear,
    fc3: Linear,
}

fn MLP::new() -> MLP {
    return MLP {
        fc1: Linear::new(784, 256),
        fc2: Linear::new(256, 128),
        fc3: Linear::new(128, 10),
    }
}

fn MLP::forward(self, x: Tensor<f32>) -> Tensor<f32> {
    let h1 = relu(self.fc1.forward(x))
    let h2 = relu(self.fc2.forward(h1))
    return self.fc3.forward(h2)
}

The Model Trait

trait Model {
    fn forward(self, input: Tensor<f32>) -> Tensor<f32>
    fn parameters(self) -> [Parameter]
    fn to(self, device: Device) -> Self
}

Device Placement

let model = MLP::new()

// Move to GPU (when available)
let gpu_model = model.to(Device::GPU(0))

// Or explicitly stay on CPU
let cpu_model = model.to(Device::CPU)