WKWesley Kambaleinkambale.dev·Feb 23 · 11 min readBuilding Neural Networks with Flax NNXOver the past two weeks, we've learned that JAX is fast (jit), that it eliminates loops (vmap), and that it computes gradients automatically (grad). These are powerful primitives. But if you've been f00
