Building Micrograd from Scratch in Rust

What is Micrograd?

Micrograd is Andrej Karpathy’s minimalist autograd engine. It implements backpropagation over a dynamically built computation graph: the same algorithm used in larger deep learning frameworks, just reduced to scalar values.

The Python version is small enough to read in one sitting. I wanted to build the same idea in Rust, where the graph has to be explicit and shared nodes need to be handled carefully.

Why Rust?

• Ownership: the graph is made of nodes that point back to their parents, so shared ownership matters.
• Interior mutability: gradients are updated during backpropagation even when nodes are shared.
• Learning: implementing autograd in Rust makes the data flow harder to hand-wave.

The Value Struct

At the heart of the crate is ValueData, a scalar plus the metadata needed for backpropagation:

1
2
3
4
5
6
7
8
9

use std::cell::RefCell;
use std::rc::Rc;

pub struct ValueData {
    pub data: f32,
    pub op: Operation,
    pub gradient: f32,
    pub parents: Vec<Value>,
}

The public handle is Value, which wraps Rc<RefCell<ValueData>>:

1
2

#[derive(Debug, Clone)]
pub struct Value(Rc<RefCell<ValueData>>);

Rc lets multiple nodes keep references to the same parent. RefCell lets the code update gradients later, during the backward pass. That is the main Rust-specific part of this implementation.

Forward Pass

Creating a leaf node stores the number and marks the operation as None:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

impl Value {
    pub fn new(data: f32) -> Value {
        Value(Rc::new(RefCell::new(ValueData {
            data,
            op: Operation::None,
            gradient: 0.0,
            parents: vec![],
        })))
    }
}

Each operation creates a new Value and stores its parents. Addition is direct: the local derivative with respect to both inputs is 1.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

impl<'a, 'b> Add<&'b Value> for &'a Value {
    type Output = Value;

    fn add(self, other: &'b Value) -> Value {
        let data = self.data() + other.data();
        Value(Rc::new(RefCell::new(ValueData {
            data,
            op: Operation::Add,
            gradient: 0.0,
            parents: vec![self.clone(), other.clone()],
        })))
    }
}

Multiplication is the same pattern, but the backward rule needs the other input’s value:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

impl<'a, 'b> Mul<&'b Value> for &'a Value {
    type Output = Value;

    fn mul(self, rhs: &'b Value) -> Self::Output {
        let data = self.data() * rhs.data();
        Value(Rc::new(RefCell::new(ValueData {
            data,
            op: Operation::Mul,
            gradient: 0.0,
            parents: vec![self.clone(), rhs.clone()],
        })))
    }
}

Backward Pass

The backward pass computes gradients using the chain rule:

The idea is simple: every node knows how it was created, so it knows how to pass its gradient to its parents.

For this expression:

Here is the backward pass one step at a time:

d is the final output, so its gradient starts at 1.0:

Then the graph is walked backwards. Since d = e + c, the gradient from d flows unchanged into both e and c:

Now $e = a \cdot b$, so multiplication sends the other input to each parent. a receives b, and b receives a:

And c = a + b, so both parents receive 1.

The important detail is that gradients add up. a is used in both e and c, so its final gradient gets one contribution from the multiply node and one from the add node.

For each node, backward applies the local derivative and adds into each parent’s gradient:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

pub fn backward(&mut self) {
    match self.op() {
        Operation::Add => {
            for parent in &self.0.borrow().parents {
                parent.0.borrow_mut().gradient += self.gradient() * 1.0;
            }
        }
        Operation::Mul => {
            let data = self.0.borrow();
            let p1 = &data.parents[0];
            let p2 = &data.parents[1];
            let p1_data = p1.data();
            let p2_data = p2.data();
            p1.0.borrow_mut().gradient += self.gradient() * p2_data;
            p2.0.borrow_mut().gradient += self.gradient() * p1_data;
        }
        Operation::None => {}
    }
}

For a multiply node:

So:

and:

That is why the code adds self.gradient() * p2_data to the first parent and self.gradient() * p1_data to the second parent.

Topological Sort

Before running backward, the graph needs to be ordered from the final output back through its dependencies. The crate does this with topo_sort:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

pub fn topo_sort(last: Value) -> Vec<Value> {
    let mut sorted: Vec<Value> = Vec::new();
    let mut visited: Vec<Value> = vec![last.clone()];
    let mut stack: Vec<Value> = vec![last.clone()];

    while !stack.is_empty() {
        if let Some(last) = stack.last().cloned() {
            let mut unvisited_flag = false;
            for parent in &last.0.borrow().parents {
                if !visited.contains(parent) {
                    visited.push(parent.clone());
                    stack.push(parent.clone());
                    unvisited_flag = true;
                }
            }
            if !unvisited_flag && let Some(last_item) = stack.pop() {
                sorted.push(last_item);
            }
        }
    }

    sorted.reverse();
    sorted[0].0.borrow_mut().gradient = 1.0;
    sorted
}

The output node starts with a gradient of 1.0, because:

Then back_propogate walks the sorted list and calls backward on each node:

1
2
3
4
5

pub fn back_propogate(last: Value) -> Vec<Value> {
    let mut list = topo_sort(last);
    list.iter_mut().for_each(|node| node.backward());
    list
}

The function name is currently spelled back_propogate in the crate, so the post uses the same spelling.

Example

Here is the example from main.rs:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

use microgradrs::{Value, back_propogate};

fn main() {
    let a = Value::new(1.0);
    let b = Value::new(2.0);
    let c = &a + &b;
    let d = &(&a * &b) + &c;
    let result = back_propogate(d);
    println!("{:#?}", result);
}

The expression is:

With a = 1 and b = 2, the gradients are:

What I Learned

1. Rc<RefCell<T>> fits this small graph well because nodes need shared ownership and delayed mutation.
2. Pointer equality matters. PartialEq uses Rc::ptr_eq, so two handles are equal only when they point to the same node.
3. The topological order is the part that makes the backward pass work. Without it, a node can be processed before its gradient is complete.

Next Steps

• Add more operations like subtraction, division, powers, relu, and tanh.
• Add tests for gradients, not just forward values.
• Build a tiny neural network layer on top of Value.

This is based on my microgradrs project.