# Grab newest JAX version.
!pip install --upgrade -q jax==0.1.54 jaxlib==0.1.37

# Make sure the Colab Runtime is set to Accelerator: TPU.
import requests
import os
if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'
  resp = requests.post(url)
  TPU_DRIVER_MODE = 1

# The following is required to use TPU Driver as JAX's backend.
from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)

import jax
import jax.numpy as np
from jax import random

key = random.PRNGKey(0)

key, subkey = random.split(key)
x = random.normal(key, (5000, 5000))

print(x.shape)
print(x.dtype)

y = np.dot(x, x)
print(y[0, 0])

x

import matplotlib.pyplot as plt

plt.plot(x[0])

np.dot(x, x.T)

print(np.dot(x, 2 * x)[[0, 2, 1, 0], ..., None, ::-1])

import numpy as onp

x_cpu = onp.array(x)
%timeit -n 1 -r 1 onp.dot(x_cpu, x_cpu)

%timeit -n 5 -r 5 np.dot(x, x).block_until_ready()

from jax import grad

def f(x):
  if x > 0:
    return 2 * x ** 3
  else:
    return 3 * x

key = random.PRNGKey(0)
x = random.normal(key, ())

print(grad(f)(x))
print(grad(f)(-x))

print(grad(grad(f))(-x))
print(grad(grad(grad(f)))(-x))

def predict(params, inputs):
  for W, b in params:
    outputs = np.dot(inputs, W) + b
    inputs = np.tanh(outputs)
  return outputs

def loss(params, batch):
  inputs, targets = batch
  predictions = predict(params, inputs)
  return np.sum((predictions - targets)**2)



def init_layer(key, n_in, n_out):
  k1, k2 = random.split(key)
  W = random.normal(k1, (n_in, n_out))
  b = random.normal(k2, (n_out,))
  return W, b

layer_sizes = [5, 2, 3]

key = random.PRNGKey(0)
key, *keys = random.split(key, len(layer_sizes))
params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))

key, *keys = random.split(key, 3)
inputs = random.normal(keys[0], (8, 5))
targets = random.normal(keys[1], (8, 3))
batch = (inputs, targets)

print(loss(params, batch))

step_size = 1e-2

for _ in range(20):
  grads = grad(loss)(params, batch)
  params = [(W - step_size * dW, b - step_size * db)
            for (W, b), (dW, db) in zip(params, grads)]

print(loss(params, batch))

from jax import jit

key = random.PRNGKey(0)
x = random.normal(key, (5000, 5000))

def f(x):
  y = x
  for _ in range(10):
    y = y - 0.1 * y + 3.
  return y[:100, :100]

f(x)

g = jit(f)
g(x)

%timeit f(x).block_until_ready()

%timeit g(x).block_until_ready()

grad(jit(grad(jit(grad(np.tanh)))))(1.0)

def f(x):
  if x > 0:
    return 2 * x ** 2
  else:
    return 3 * x

g = jit(f)

f(2)

g(2)

def f(x, n):
  i = 0
  while i < n:
    x = x * x
    i += 1
  return x

g = jit(f)

f(np.array([1., 2., 3.]), 5)

g(np.array([1., 2., 3.]), 5)

g = jit(f, static_argnums=(1,))

g(np.array([1., 2., 3.]), 5)

from jax import vmap

print(vmap(lambda x: x**2)(np.arange(8)))

from jax import make_jaxpr

make_jaxpr(np.dot)(np.ones(8), np.ones(8))

make_jaxpr(vmap(np.dot))(np.ones((10, 8)), np.ones((10, 8)))

make_jaxpr(vmap(vmap(np.dot)))(np.ones((10, 10, 8)), np.ones((10, 10, 8)))

perex_grads = vmap(grad(loss), in_axes=(None, 0))
make_jaxpr(perex_grads)(params, batch)

jax.devices()

from jax import pmap

y = pmap(lambda x: x ** 2)(np.arange(8))
print(y)

y

z = y / 2
print(z)

import matplotlib.pyplot as plt
plt.plot(y)

keys = random.split(random.PRNGKey(0), 8)
mats = pmap(lambda key: random.normal(key, (5000, 5000)))(keys)
result = pmap(np.dot)(mats, mats)
print(pmap(np.mean)(result))

timeit -n 5 -r 5 pmap(np.dot)(mats, mats).block_until_ready()

from functools import partial
from jax.lax import psum

@partial(pmap, axis_name='i')
def normalize(x):
  return x / psum(x, 'i')

print(normalize(np.arange(8.)))

@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def f(x):
  row_sum = psum(x, 'rows')
  col_sum = psum(x, 'cols')
  total_sum = psum(x, ('rows', 'cols'))
  return row_sum, col_sum, total_sum

x = np.arange(8.).reshape((4, 2))
a, b, c = f(x)

print("input:\n", x)
print("row sum:\n", a)
print("col sum:\n", b)
print("total sum:\n", c)

@pmap
def f(x):
  y = np.sin(x)
  @pmap
  def g(z):
    return np.cos(z) * np.tan(y.sum()) * np.tanh(x).sum()
  return grad(lambda w: np.sum(g(w)))(x)
  
f(x)

grad(lambda x: np.sum(f(x)))(x)

from jax import jvp, vjp  # forward and reverse-mode

curry = lambda f: partial(partial, f)

@curry
def jacfwd(fun, x):
  pushfwd = partial(jvp, fun, (x,))  # jvp!
  std_basis = np.eye(onp.size(x)).reshape((-1,) + np.shape(x)),
  y, jac_flat = vmap(pushfwd, out_axes=(None, -1))(std_basis)  # vmap!
  return jac_flat.reshape(np.shape(y) + np.shape(x))

@curry
def jacrev(fun, x):
  y, pullback = vjp(fun, x)  # vjp!
  std_basis = np.eye(onp.size(y)).reshape((-1,) + np.shape(y))
  jac_flat, = vmap(pullback)(std_basis)  # vmap!
  return jac_flat.reshape(np.shape(y) + np.shape(x))

def hessian(fun):
  return jit(jacfwd(jacrev(fun)))  # jit!

input_hess = hessian(lambda inputs: loss(params, (inputs, targets)))
per_example_hess = pmap(input_hess)  # pmap!
per_example_hess(inputs)