JAX from NumPy to Machine Learning

Move from familiar NumPy-style Python to the powerful world of JAX and modern machine learning, one practical step at a time.

JAX combines the familiar array-based programming style of NumPy with powerful capabilities for automatic differentiation, JIT compilation, vectorization, and accelerated numerical computing. But for beginners, learning how these pieces fit together can feel overwhelming.

JAX from NumPy to Machine Learning provides a clear, hands-on path from basic JAX arrays to building and training a working neural network.

Rather than treating JAX as a collection of disconnected features, this book shows you how its core ideas work together. You will begin with familiar numerical operations, gradually learn the programming patterns that make JAX different, and then apply those skills to a complete machine-learning workflow.

Inside, you will learn how to:

• Work confidently with JAX arrays, shapes, data types, broadcasting, indexing, and immutable updates

• Move existing NumPy-style calculations into JAX

• Understand and compute gradients with automatic differentiation

• Use jax.grad() and value_and_grad() in practical optimization problems

• Speed up repeated computations with jax.jit()

• Understand tracing, static values, and common causes of unnecessary recompilation

• Replace repetitive batch-processing code with jax.vmap()

• Combine grad, jit, and vmap effectively

• Organize model parameters with PyTrees

• Build a neural network from JAX functions

• Implement forward passes, loss calculations, and accuracy measurements

• Train a model with gradient descent and mini-batch training

• Build a JIT-compiled training step

• Evaluate predictions and diagnose common training problems

• Understand JAX devices, data placement, benchmarking, and performance issues

• Refactor your work into a cleaner, more efficient JAX project

The book keeps the mathematics approachable and focuses on understanding what the code is doing, why each JAX feature matters, and when to use it.

You do not need previous experience with JAX, deep-learning frameworks, GPU programming, or advanced mathematics. Basic Python knowledge is enough to begin, and familiarity with NumPy is helpful but not required.

Whether you are coming from Python, NumPy, scientific computing, or introductory machine learning, JAX from NumPy to Machine Learning will help you build the practical foundation you need to start writing efficient, transformation-friendly JAX programs with confidence.

Start with arrays. Understand gradients. Compile and vectorize your code. Then bring everything together by building and training your own neural network with JAX.