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.…