JAX is a Python library that has a very numpy-like interface and enables running numerically intensive code on a variety of devices including CPUs, GPUs, and TPUs. Built in automatic differentiation makes JAX a great tool for writing machine learning code. There’s a lot of good technical reasons to learn jax including the fact that it is almost interchangeable with numpy, its support for just in time (JIT) compilation, auto-vectorization with vmap() and its ability to easily scale workflows to a large number of devices.
Why learn machine learning with JAX
There are many more technical reasons to learn JAX, such as its functional programming philosophy, its approach to randomness and replicability, etc. However, what I want to focus on here is the fact that JAX is a great tool for learning machine learning.
Simplicity is at the heart of JAX as a framework. The state of your code and models is almost always in front of you. Details are rarely abstracted away. Granted, this can be a matter of preference, but I genuinely like the fact the library doesn’t try to hide the gory details from you.
How the “book” came about
You can find the latest version of the “book” here.
Please note that for the time being this is a work in progress. I’m trying to weaponize the embarrassment of having something incomplete out there into motivation to finish it.
A huge part of what got me started on writing the “book” is that I spent some time designing a graduate course in machine learning for life sciences students (Hola jeffe!). Both Pytorch and Tensorflow/Keras make excellent tools and have a mature ecosystem. However, the problem with having a mature ecosystem is that frameworks are starting to become opinionated. While there are good reasons for that, this is not always a good thing, especially if you’re trying to teach the basics of machine learning (don’t blame me, I did most of my machine learning phd in matlab with the occasional sprinkles of c++ for performance reasons).
Numpy was an obvious choice for illustrating the fundamentals of machine learning, but it lacks a lot of the niceties of deep learning frameworks. Most notable is the lack of support for accelerators like GPUs, as well as the lack of automatic differentiation. This is where JAX comes in. JAX has a very numpy-like interface and enables running numerically intensive code on a variety of devices including CPUs, GPUs, and TPUs, and built-in autodiff makes JAX a great tool for writing machine learning code.
About that machine learning part
Machine learning theory is beautiful but often skipped in deep learning courses. Partly this is because deep learning courses assume that the audience is familiar with the basics of machine learning.
A common pattern I’ve noticed from professional experience in interview settings is that deep learning candidates often have gaps in what I would think are fundamentals. For example they would struggle to explain how ideas like the bias-variance decomposition affect their work on a practical level, or tell me why they might prefer a simpler linear model over a deep learning model. Another common pattern I’ve seen is candidates who have a very good grasp of the theory and practice of developing deep learning models, but fail to extend this knowledge on the systems-level aspects of machine learning. In the real world, datasets don’t appear out of thin air. People need to design data collection mechanisms. The data will then need to be cleaned, and preprocessed. Models need to be trained, evaluated, and deployed. Model deployments need to be monitored, and updated as new data becomes available. There is a lifecycle at work here.
In designing this course I’m aiming to cover some of the basics before going to the actual deep learning topics. I will also be trying to give ideas as to how the lifecycle of machine learning systems works out in the real world whenever possible.
What you will learn
- The basics of machine learning and the machine learning lifecycle
- How to use assess machine learning models
- How to use JAX to write machine learning code
- Methods and tools to optimize your machine learning pipelines
- Examples of where and how machine learning is used in the real world
Who is this for?
- People who want a broad overview of machine learning
- People who want to learn JAX
- People who want to learn how to write numerical code
- People who want to learn how to write code that runs on a variety of devices
What you will need
- A computer with a modern browser
- A reliable internet connection
- Some basic knowledge of Python is going to be very helpful
About me
I am machine learning researcher working on understanding the dynamics of the open source ecosystem at Quine. Prior to Quine, I worked on in-browser privacy preserving machine learning for Brave, and built machine learning systems in a number of domains, including biotech, life sciences, finance and security. My interests include machine learning, health & biomedicine, natural language processing, and making GPUs overheat, a habbit for which I occasionally get published in peer-reviewed venues. For more information about me visit my peronal website at gotbim.com.
—>