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: /
JAX
SOCIAL SHARE CARD GENERATOR