From NumPy to JAX: My First "Aha!" Moments with Accelerated AI

Devansh Bajaj Posted on May 30 From NumPy to JAX: My First "Aha!" Moments with Accelerated AI # machinelearning # ai # python # tutorial Building open-source solutions for my 100 Days of AI Agents challenge meant I needed to start looking at frameworks that scale better than standard NumPy and PyTorch. That inevitably led me to JAX. Transitioning to JAX requires a bit of a paradigm shift. If you are used to the standard Python data science stack, JAX forces you to rewire how you think about array operations, memory, and hardware execution. I spent today digging into the core mechanics, and I want to share my top 3 takeaways and the exact code snippets that made it click for me. 1. Immutability is a Feature, Not a Bug This was my first major roadblock. In standard NumPy, if you want to change an element in an array, you just reassign it in place. Python import numpy as np x = np.arange(10) x[0] = 10 print(x) # Output: [10 1 2 3 4 5 6 7 8 9] If you try the exact same thing in JAX, it screams at you: TypeError: JAX arrays are immutable. JAX arrays (jax.Array) cannot be changed once created. This is a core design principle that enables JAX's functional programming nature and automatic differentiation. To update an array, JAX provides an indexed update syntax that returns an updated copy: Python **import jax.numpy as jnp x = jnp.arange(10) y = x.at[0].set(10) print(y) # Output: [10 1 2 3 4 5 6 7 8 9] print(x) # Output: 0 1 2 3 4 5 6 7 8 9 .** The Catch: This does create memory overhead since you are creating copies, but it completely eliminates the side-effects that make distributed computing a nightmare. 2. Native Hardware Awareness & Sharding JAX arrays inherently know where they live. You don't have to jump through hoops to figure out if your data is on the CPU, GPU, or TPU. By default, JAX pushes operations to the fastest available accelerator. Running this locally on my MSI Raider, I can easily inspect exactly where my array is stored using .devices():