mlframework

Building a rudimentary ml framework from scratch in rust.

Note

This project is a work in progess. It is also intended as a learning project to improve my understanding of ML framework implementation details and Rust.

Features

Basic operations with backprop

let x = Tensor::new([3., 4., 5.]);
let y = Tensor::new([1., -2., 1.]);
let z = Tensor::new([-3., 1., 3.]);
let s = ((x + y.clone() + y) * z.relu()).reduce_sum();

// Performs reverse mode autodiff and sets private grad field on tensors throughout the graph of `s`.
s.backward();

Compile-time shape checking

let x = Tensor::new([[0.0; 5]; 10]); // shape: [10, 5]

// Simply specify a Tensory w/ dtype and shape and call reshape
// If the shape is valid for the starting shape the code will compile
let x2: Tensor<f32, D3<2, 5, 5>> = x.reshape(); // shape: [2, 5, 5]

// This will *not* compile, becase the shape [51] is invalid for a tensor of shape [2, 5, 5]
let x3: Tensor<f32, D1<51>> = x2.reshape(); // ERROR!!

Status

• Initial tensor structure
• Basic elementwise operations
• Basic backprop implementation
  • Fix graph traversal complexity
• Tensor Shapes
• Matmul
• Tensor Indexing
• Potentially improve Op implementation
  • Make it simpler to create/register new ops & backward functions
  • Generate similar ops using macros
• Define a module structure
• Implement an optimizer
• Setup test structure
• Organize repo code
• Add CI w/ Github Actions