High Level Neural Network Modeling with FlaxIn our previous article, we broke down why JAX requires explicit state handling. We saw that wrapping our parameters, optimizer momentum, and PRNG keys into a clean carry pattern lets us use jax.jit, 71JAug 19 · 7 min readmarshud.dev
Functional State Management and Low-Level Control Flow in JAXWell be back guys, in our previous article on managing PRNG keys, we explored how JAX abandons the conventional hidden random state in favor of stateless random number generation. We learned that to g1Aug 11 · 11 min readmarshud.dev
Managing PRNG Keys in JAXIn the previous articles, we have explored the fundamental principles of JAX from functional purity to transformations to how it handles arrays to PyTrees. However, in real-world ML, when working withAug 6 · 7 min readmarshud.dev
Program Transformations and PyTrees in JAXWelcome back everyone, in our previous article, we covered the foundational philosophy of JAX from functional purity to array immutability to basic transformations like jax.grad, jax.vmap and jax.jit.Aug 4 · 9 min readmarshud.dev
JAX Core Fundamentals: Pure Functional ArraysWelcome everyone, in this article, I will explore the core philosophy of JAX, it's array system and the critical concept of functional purity. By the end of this article, you will understand why JAX w5Jul 7 · 11 min readmarshud.dev
