STAT 39000: Project 11 — Spring 2022

JAX is a library for high performance computing. It falls into the same category as other popular packages like: numpy, pytorch, and tensorflow. JAX is a product of Google / Deepmind that takes a completely different approach than their other product, tensorflow.

Like the the other popular libraries, JAX can utilize GPUs/TPUs to greatly speed up computation. Let’s take a look.

Here is a snippet of code from previous projects that uses pytorch and calculates predictions 10000 times.

Approximately how much time does it take to run this second chunk of code (after we have already read in our data)?

Here is the equivalent JAX code:

How much time does this take?

At this point in time you may be questioning how JAX could possibly be worth it. At first glance, the new code does look a bit cleaner, but not clean enough to use code that is around 3 times slower.

This is where JAX first trick, or transformation comes in to play. When we refer to transformation, think of it as an operation on some function that produces another function as an output.

The first transformation we will talk about is jax.jit. "JIT" stands for "Just In Time" and refers to a "Just in time" compiler. Essentially, just in time compilation is a trick that can be used to greatly speed up the execution of some code by compiling the code. In a nutshell, the compiled version of the code has a wide variety of optimizations that speed your code up.

Lots of our computation time is spent inside our loop, specifically when we are calculating our y_predictions. Let’s see if we can use the jit transformation to speed up our JAX code with little to no extra effort.

Write a function called model that accepts two arguments. The first argument is a tuple containing our parameters: beta0, beta1, and beta2. The second is our input to our function (our x values) called x. model should then unpack our tuple of parameters into beta0, beta1, and beta2, and then return predictions (the same formula shown above, twice). Replace the code as follows.

Run and time the code again. No difference? Well, we didn’t use our jit transformation yet! Using the transformation is easy. JAX provides two equivalent ways. You can either decorate your model function with the @jax.jit decorator, or simply apply the transformation to your function and save the new, jit compiled function and use it instead.

Re-run your code using the JIT transformation. Is it faster now?

At this point in time you may be considering slapping @jax.jit on all your functions — unfortunately it is not quite so simple! First of all, the previous comparison was actually not fair at all. Why? JAX has asynchronous dispatch by default. What this means is that, by default, JAX will return control to Python as soon as possible, even if it is before the function has been fully evaluated.

What does this mean? It means that our finished example from question 1 may be returning a not-yet-complete result, greatly throwing off our performance measurements. So how can we synchronously wait for execution to finish? This is easy, simply use the block_until_ready method built in to your jit compiled model function.

Re-run your code from before — you should find that the results are unchanged, it turns out that really was a serious speedup from before. Great. Let’s move on from this part of things. Back to our question. Why can’t we just slap @jax.jit on any function and expect a speedup?

Take the following function.

If you try running it you will get an error saying something along the lines of "TracerIntegerConversionError". The problem with this function, and why it cannot be jit compiled, is the epochs argument. By default, JAX tries to "trace" the parameters to determine its effect on inputs of a specific shape and type. Control flow cannot depend on traced values — in this case, epochs is relied on in order to determine how many times to loop. In addition, the shapes of all input and output values of a function must be able to be determined ahead of time.

How do we fix this? Well, it is not always possible, however, we can choose to select parameters to be static or not traced. If a parameter is marked as static, or not traced, it can be JIT compiled. The catch is that any time a call to the function is made and the value of the static parameter is changed, the function will have to be recompiled with that new static value. So, this is only useful if you will only occasionally change the parameter. This sounds like our case! We only want to occasionally change the number of epochs, so perfect.

You can mark a parameter as static by specifying the argument position using the static_argnums argument to jax.jit, or by specifying the argument name using the static_argnames argument to jax.jit.

Force the epochs argument to be static, and use the jax.jit decorator to compile the function. Test out the function, in order using the following code cells.

Do your best to explain why the last code cell was once again slower.

We learned that one of the coolest parts of the pytorch package was the automatic differentiation feature. It saves a lot of time doing some calculus and coding up resulting equations. Recall that in pytorch this differentiation was baked into the backward method of our MSE. This is quite different from the way we think about the equations when looking at the math, and is quite confusing.

JAX has the same functionality, but it is much cleaner and easier to use. We will provide you with a simple example, and explain the math as we go along.

Let’s say our function is $f(x) = 2x^2$. We can start by writing a function.

Fantastic, so far pretty easy.

The derivative w.r.t. x is $4x$. Doing this in JAX is as easy as applying the jax.grad transformation to the function.

Okay, test out both functions as follows.

On the last line, you probably received a message or error saying something along the lines of "Gradient only defined for scalar-ouput functions. What this means is that the resulting derivative function is not vectorized. As you may have guessed, this is easily fixed. Another key transformation that JAX provides is called vmap. vmap takes a function and creates a vectorized version of the function. See the following.

Heck yes! That is pretty cool, and very powerful. It is so much more understandable than the magic happening in the pytorch world too!

Dig back into your memory about any equation you may have had in the past where you needed to find a derivative. Create a Python function, find the derivative, and test it out on both a single value, like 4.0 as well as an array, like jnp.array([1.0,2.0,3.0]). Don’t hesitate to make it extra fun and include some functions like jnp.cos, jnp.sin, etc. Did everything work as expected?

Okay, great, but that was a straight-forward example. What if we have multiple parameters we’d like to take partial derivatives with respect to? jax.grad can handle that too!

Read this excellent example in the official JAX documentation.

Given the following (should be familiar) model, create a function called get_partials that accepts an argument params (a tuple containing beta0, beta1, and beta2, in order) and an argument x, that can be either a single value (a scalar), or a jnp.array with multiple values. This function should return a single value for each of the 3 partial derivatives, where x is plugged into each of the 3 partial derivatives to calculate each value, OR, 3 arrays of results where there are 3 values for each value in the input array.

Since we have 3 parameters, we will have 3 partial derivatives, and our new function should output a value for each of our 3 partial derivatives, for each value passed as x. To be explicit and allow you to check your work, the results should be the same as the following.