JAX¶

Christian Elsasser (che@physik.uzh.ch)

The content of this lecture might be distributed under CC by-sa.

In [1]:
import jax # JAX functionality
import jax.numpy as jnp # Numpy array of JAX
import numpy as np # Regular numpy array

We create a large numpy array

In [2]:
a = np.random.randn(10_000,10_000)
In [3]:
j = jnp.array(a)

We create a third-order polynomial function and its (jit-)compiled version

In [4]:
def f(x):
    return -4*x*x*x + 9*x*x + 6*x - 3

@jax.jit # <- this compiles the function 
def f_compiled(x):
    return -4*x*x*x + 9*x*x + 6*x - 3
# we can also do f_c = jax.jit(f)

Let's benchmark the different cases

In [5]:
%timeit f(a)                                         # NumPy
%timeit f(jnp.array(a)).block_until_ready()          # JAX
%timeit f_compiled(jnp.array(a)).block_until_ready() # JAX + JIT pre-compiled
%timeit jax.jit(f)(jnp.array(a)).block_until_ready() # JAX + JIT 

# Since JAX is operating asynchronously we need to ensure that the command is
# conclude before the next loop is started.

1.8 s ± 14.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
874 ms ± 5.64 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
287 ms ± 8.35 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
281 ms ± 4.54 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

We assumed that for the JAX cases the jnp.array has to be only created for this operation which is rather conservative and disfavourably for JAX.
So let's test it with the array already created

In [6]:
j = jnp.array(a)
In [7]:
%timeit f(j).block_until_ready()          # JAX
%timeit f_compiled(j).block_until_ready() # JAX + JIT pre-compiled
%timeit jax.jit(f)(j).block_until_ready() # JAX + JIT 

665 ms ± 2.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
66.3 ms ± 347 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
66.6 ms ± 250 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [8]:
%timeit np.exp(j)
%timeit jnp.exp(j).block_until_ready()

153 ms ± 2.76 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
59.3 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [9]:
%timeit (2*a+3)
%timeit (2*j+3).block_until_ready()

465 ms ± 2.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
127 ms ± 606 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Auto-differentiation¶

In [ ]:
import numdifftools as nd
In [ ]:
# Analytical
df_a = lambda x : -12*x*x + 18*x + 6
df_a_comp = jax.jit(df_a)

# Numerical
df_n = nd.Derivative(f)

# Auto-differentiation
df_j = jax.grad(f)
df_j_comp = jax.jit(df_j)
In [ ]:
b = np.random.randn(1_000,1_00)
In [ ]:
%timeit df_a(b) # Analytical derivative
%timeit df_a_comp(b) # Analytical derivative compiled

%timeit df_n(b) # Numerical derivative

%timeit jnp.vectorize(df_j)(jnp.array(b)).block_until_ready() # Auto-diff
%timeit jnp.vectorize(df_j_comp)(jnp.array(b)).block_until_ready() # Auto-diff compiled
In [ ]: