JAX on CoCalc

Kernel: Python 3 (Ubuntu Linux)

JAX is Autograd and XLA, brought together for high-performance machine learning research

grad for gradient, and jit for just in time compilation

jax also has a numpy compatible interface

In [1]:

from jax import grad, jit
import jax.numpy as np

In [2]:

def tanh(x):  # Define a function
    y = np.exp(-2.0 * x)
    return (1.0 - y) / (1.0 + y)

In [3]:

grad_tanh = grad(tanh)  # Obtain its gradient function
print(grad_tanh(1.0))

In [4]:

%timeit grad_tanh(0.1)

In [5]:

def slow_f(x):
    # Element-wise ops see a large benefit from fusion
    a = x * x
    b = 2.0 + x
    return a * b

In [6]:

fast_f = jit(slow_f)

In [7]:

x = np.ones((5000, 5000))

In [8]:

%timeit -n10 -r3 fast_f(x)

In [9]:

%timeit -n10 -r3 slow_f(x)

In [10]:

del x

comparing with numba

In [11]:

import numba
numba_f = numba.jit(slow_f)

In [12]:

import numpy
y = numpy.ones((5000, 5000))

In [13]:

%timeit -n10 -r3 numba_f(y)

In [0]:

In [0]: