pc-rs

Predictive coding in Rust.

mod gaussian;
mod graph;
mod linalg;

use gaussian::{function::GaussianFunction, variable::GaussianVariable};
use graph::Graph;
use linalg::{math::Activation, matrix::Matrix, vector::Vector};

const MU_SIZE: usize = 10;
const DATA_A_SIZE: usize = 128;
const DATA_B_SIZE: usize = 128;
const BATCH_SIZE: usize = 64;
const NUM_EPOCHS: usize = 10;
const NUM_ITERATIONS: usize = 12;

fn main() {
    let mu = GaussianVariable::new(Matrix::normal(BATCH_SIZE, MU_SIZE, 0.0, 0.05), false);
    let data_a = GaussianVariable::new(Matrix::ones(BATCH_SIZE, DATA_A_SIZE) * 2.0, true);
    let data_b = GaussianVariable::new(Matrix::ones(BATCH_SIZE, DATA_B_SIZE) * 4.0, true);

    let mu_data_a = GaussianFunction::new(
        Matrix::normal(MU_SIZE, DATA_A_SIZE, 0.0, 0.05),
        Vector::zeros(DATA_A_SIZE),
        Activation::Linear,
    );
    let mu_data_b = GaussianFunction::new(
        Matrix::normal(MU_SIZE, DATA_B_SIZE, 0.0, 0.05),
        Vector::zeros(DATA_B_SIZE),
        Activation::Linear,
    );

    let mut graph = Graph::<GaussianVariable, GaussianFunction>::new();
    let mu_index = graph.add_node(mu);
    let data_a_index = graph.add_node(data_a);
    let data_b_index = graph.add_node(data_b);

    graph.add_edges(vec![
        (mu_index, data_a_index, mu_data_a),
        (mu_index, data_b_index, mu_data_b),
    ]);

    for _ in 0..NUM_EPOCHS {
        let mu = graph.get_node_mut(mu_index).unwrap();
        mu.set_data(Matrix::normal(BATCH_SIZE, MU_SIZE, 0.0, 0.05));

        for _ in 0..NUM_ITERATIONS {
            graph.infer();
        }
        graph.learn();
    }

    let preds = graph.forward();
    let mu_data_a_pred = preds.get(&(mu_index, data_a_index)).unwrap();
    let mu_data_b_pred = preds.get(&(mu_index, data_b_index)).unwrap();
    println!("mu_data_a_pred: {:?}", mu_data_a_pred.mean());
    println!("mu_data_b_pred: {:?}", mu_data_b_pred.mean());
}

To run a demo:

cargo run

To run tests:

cargo test

Tasks

Create Variable trait that T and F implement
Create optim which accept Gradient objects and calls their update function, uses index to reference to gradient and update the graph
graph.infer adds gradients to optim
optim.step updates gradients and updates the graph
Setup an optim for both variables and functions

pub trait Derivative {}

pub trait Variable {
    fn update(&mut self, derivatives: &Vec<Gradient>);
}

pub trait Function<T: Variable, G: Gradient, D: Gradient> {
    fn forward(&self, input: &T) -> Matrix;
    fn backward(&self, input: &T, target: &T) -> (G, G);
    fn backward_params(&self, input: &T, target: &T) -> D;
}

alec-tschantz/pc-rs

pc-rs

Tasks