Repository files navigation

mojograd is a Mojo implementation ofmicrograd,a reverse-mode autodiff library with a PyTorch-like API.

The goal is to be as close as possible to micrograd, keeping apretty clean syntax to define computational graphs.Like micrograd, it only supports scalar values for now, but weplan to extend it to support Tensors in the near future.

Note thatmojograd is in WIP and relies on static register passable structures, sobackward pass copies values and can be really slow (Mojo traits support shouldimprove that, so please stay tuned!). However, even now with zerooptimizations, forward pass is already ~40x faster than the original Python implementation(see benchmarks bellow).

mojograd dynamically builds a computational graph by overloading operatorsonValue type, performing the forward pass. Just write your expression likea normal (non-diff) equation and callbackward() to perform thebackward pass:

frommojogradimportValuevara=Value(2.0)varb=Value(3.0)varc:Float32=2.0vard=b**cvare=a+ce.backward()a.print()# => <Value data: 2.0 grad: 1.0 op:  >b.print()# => <Value data: 3.0 grad: 0.0 op:  >d.print()# => <Value data: 9.0 grad: 0.0 op: ** >e.print()# => <Value data: 4.0 grad: 1.0 op: + >

For a more complete example (a simple Multi-Layer Perceptron), please checkthetests.mojo file. You can run it with:

mojo tests.mojo

When compared to original Python implementation,mojograd is up to~40 times fasterin forward pass.

• 2023.11.19
  • Benchmarking inference and comparing with micrograd
• 2023.11.18
  • Optimization pass through the code
• 2023.11.14
  • Rebuild the whole thing using pointer handling (dangerous) to register-passables
  • Got the full micrograd implementation working!
  • MLP example training and inference working!
• 2023.09.05
  • Starting from scratch based on suggestions from Jack Clayton
  • Topological sort works but I'm messing something with memory handling,the gradients are not getting updated
• 2023.07.04
  • Ported Neuron, Layer and MLP
  • Back to use yakupc55's List (needregister_passable data struct)
• 2023.06.30
  • Finally got it working! Only missing pow ops and review it