Modern Machine Learning Software

Scientific machine learning turns mathematical models into programs that must be differentiated, optimized, and evaluated efficiently on modern hardware. Ordinary Python makes experimentation convenient, but these transformations require a disciplined representation of functions, data, and model parameters.

The treatment assumes scientific Python, linear algebra, multivariable calculus, and the basic machine learning concepts stated in the book prerequisites. JAX, functional programming, and accelerator programming are introduced from first principles.

We begin with pure functions, compilation, vectorization, and explicit random-number state. We then use type signatures and pytrees to describe structured computations, develop automatic differentiation from computational graphs, and assemble these tools into optimization algorithms for model training. The resulting software foundation supports concise implementations that JAX can transform, compile, and run on central processing units (CPUs) and graphics processing units (GPUs).