Custom training loop from scratch in JAX
Writing a training loop that you can actually control
Mar 4, 202610 min read36
Command Palette
Search for a command to run...
#jax
Articles tagged with #jax
Building Neural Networks with Flax NNX
Flax NNX: Neural Networks That Feel Like PyTorch
Feb 23, 202611 min read124
Transformations That Change Everything
Automatic vectorization with vmap and gradients with grad in Jax
Feb 9, 202610 min read78
Why JAX? The NumPy You Know, But Faster
How JAX uses XLA to strap a rocket engine onto your math operations.
Feb 2, 20269 min read127
Building Convolutional Neural Networks in JAX
A comprehensive walkthrough of building a CNN using JAX, from data preprocessing to training and evaluation
Mar 12, 20259 min read1.0K
NumPy vs JAX: Optimizing Performance and Functionality
A beginner's guide and comparison of NumPy and JAX for scientific computing and machine learning
May 8, 20236 min read3.0K