Search for a command to run...
Series
This series will build someone from basic to a solid understanding of the JAX Eco System
Welcome 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 w
Welcome 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.
In 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 with
Well 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 g
In 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,
