Part 7/12:
Conversion from NumPy is simple—replace np with jnp—and the operations are similar. However, it is vital to note that JAX enforces immutability, meaning in-place updates are not allowed. Instead, to update arrays, one must create new copies via functions like jax.ops.index_update.
The Philosophical Backbone: Immutability and Pure Functions
JAX strongly encourages functional programming principles, which promote pure functions—functions that return consistent outputs for the same inputs without side-effects. This design simplifies debugging, ensures reproducibility, and aligns with the way accelerators process data.