What is JAX?
JAX is Autograd and XLA, brought together for high-performance numerical computing and ML research. It provides composable transformations of Python+NumPy programs: differentiate, vectorize, parallelize, JIT compile to GPU/TPU, and more.
🤔🧐
#30DaysOfJAX
1/11🧵
JAX is Autograd and XLA, brought together for high-performance numerical computing and ML research. It provides composable transformations of Python+NumPy programs: differentiate, vectorize, parallelize, JIT compile to GPU/TPU, and more.
🤔🧐
#30DaysOfJAX
1/11🧵
That's already a lot to take in!
Let's try to understand the key words first
What is:
• Autograd
• XLA
• Differentiate
• Just-in-time compile
2/11🧵
Let's try to understand the key words first
What is:
• Autograd
• XLA
• Differentiate
• Just-in-time compile
2/11🧵
What is Differentiate?
🚨 If you studied Calculus you might remember this one. (bear with me, don't run!)
Imagine you have a function:
-> f(x) = 3*x + 4
and you want to know how sensible your output (f) is to changes in the input (x).
3/11🧵
🚨 If you studied Calculus you might remember this one. (bear with me, don't run!)
Imagine you have a function:
-> f(x) = 3*x + 4
and you want to know how sensible your output (f) is to changes in the input (x).
3/11🧵
That's what we call the Derivative of f with respect of x
For our example, you can see that for a change in 1 on x, it will impact in a change of 3 on f
So the Derivative of f with respect to x is 3
Differentiation is the process to find the derivative
4/11🧵
For our example, you can see that for a change in 1 on x, it will impact in a change of 3 on f
So the Derivative of f with respect to x is 3
Differentiation is the process to find the derivative
4/11🧵
When we have a function with many different variables:
-> f(x1, x2, ...) = a1*x1 + a2*x2 + … + b
We calculate the derivative of f with respect to x1, f with respect to x2…
Each one is called a partial derivative
The vector with all of them is called gradient!
5/11🧵
-> f(x1, x2, ...) = a1*x1 + a2*x2 + … + b
We calculate the derivative of f with respect to x1, f with respect to x2…
Each one is called a partial derivative
The vector with all of them is called gradient!
5/11🧵
The Gradient is the vector with all the partial derivatives of a multivariate function
To understand the meaning, think like this:
-> If you are climbing a mountain, and you have a function (f) that tells you altitude
-> The gradient of f will point to the top of the ⛰️
6/11🧵
To understand the meaning, think like this:
-> If you are climbing a mountain, and you have a function (f) that tells you altitude
-> The gradient of f will point to the top of the ⛰️
6/11🧵
What is Autograd?
It is a library that can automatically differentiate native Python and Numpy code.
Given your def(a, b, c), you can find it's gradient
The main intended application of Autograd is gradient-based optimization
github.com/hips/autograd
7/11🧵
It is a library that can automatically differentiate native Python and Numpy code.
Given your def(a, b, c), you can find it's gradient
The main intended application of Autograd is gradient-based optimization
github.com/hips/autograd
7/11🧵
What is XLA?
Accelerated Linear Algebra (XLA) is a domain-specific compiler for linear algebra that can accelerate math operations with potentially no source code changes.
tensorflow.org/xla
8/11🧵
Accelerated Linear Algebra (XLA) is a domain-specific compiler for linear algebra that can accelerate math operations with potentially no source code changes.
tensorflow.org/xla
8/11🧵
What is JIT? ⚡️
Just-in-time compilation is technique that interpreted languages use that, while executing code, the interpreter will also compile the code so that on the next time that the code is executed, it will run as compiled code (faster) ⚡️
9/11🧵
Just-in-time compilation is technique that interpreted languages use that, while executing code, the interpreter will also compile the code so that on the next time that the code is executed, it will run as compiled code (faster) ⚡️
9/11🧵
With all the concepts (loosely explained for brevity), we can understand what is JAX:
JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research.
10/11🧵
JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research.
10/11🧵
This month I'll keep posting about my journey learning JAX. 📚👓
Now I know what it is! Next: how to use it!
Join me and let's learn together!
Follow me so you don't miss the next updates!
#30DaysOfJAX
11/11🧵
Now I know what it is! Next: how to use it!
Join me and let's learn together!
Follow me so you don't miss the next updates!
#30DaysOfJAX
11/11🧵
• • •
Missing some Tweet in this thread? You can try to
force a refresh
