MMarshudinmarshud.dev·18h ago · 11 min readFunctional 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 g00
MMarshudinmarshud.dev·6d ago · 7 min readManaging 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 with00
MMarshudinmarshud.dev·Aug 4 · 9 min readProgram 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.00
MMarshudinmarshud.dev·Jul 7 · 11 min readJAX 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 w50
MMarshudinmarshud.dev·Jul 3 · 4 min readWhy the JAX AI Stack is built for the Future of Foundation ModelsFoundation models like large-scale transformers, multimodal systems, MoEs, etc are pushing the currently known boundaries of AI. Training and serving them demands extreme scale, hardware efficiency, r20