Visualizing SGD doing its thing 🥾
Nabla + MAX + Matplotlib + NumPy: github.com/nabla-ml/nabla…
- I am starting a (notebook) series on training transformers with Nabla. Part 1 is a side-by-side (Nabla vs. JAX) toy implementation from scratch: nablaml.com/tutorials/jax_…
- Automatic Vectorization (vmap) in action: Write a program once, then use it for any batched input. If applied correctly, this can greatly reduce the number of for-loops and speed up a program. 🎓 Learn more about visualizing program transformations: nablaml.com/tutorials/unde…
- JAX vs. Nabla: An initial speed comparison (on cpu) for training an MLP on a simple regression task. 🤗 The full Notebook: nablaml.com/tutorials/nabl…
- I am reverse-engineering JAX from scratch in Python, but instead of using XLA, I am using NumPy @numpy_team and MAX @Modular for CPU/GPU acceleration. 🐍🫦 Working: Function transforms like vmap, grad, jit etc., some built-in nn/ modules, pip-install.
