# 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 [ ]: