import jax.numpy as jnp

from jax import pmap

result = pmap(lambda x: x ** 2)(jnp.arange(7))
print(result)

from jax import vmap

x = jnp.array([1., 2., 3.])
y = jnp.array([2., 4., 6.])

print(vmap(jnp.add)(x, y))
print(pmap(jnp.add)(x, y))

from jax import make_jaxpr

def f(x, y):
  a = jnp.dot(x, y)
  b = jnp.tanh(a)
  return b

xs = jnp.ones((8, 2, 3))
ys = jnp.ones((8, 3, 4))

print("f jaxpr")
print(make_jaxpr(f)(xs[0], ys[0]))

print("vmap(f) jaxpr")
print(make_jaxpr(vmap(f))(xs, ys))

print("pmap(f) jaxpr")
print(make_jaxpr(pmap(f))(xs, ys))

y = pmap(lambda x: x ** 2)(jnp.arange(8))
z = y / 2
print(z)

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

y

y / 2

import numpy as np
np.sin(y)

from jax import random

# create 8 random keys
keys = random.split(random.key(0), 8)
# create a 5000 x 6000 matrix on each device by mapping over keys
mats = pmap(lambda key: random.normal(key, (5000, 6000)))(keys)
# the stack of matrices is represented logically as a single array
mats.shape

# run a local matmul on each device in parallel (no data transfer)
result = pmap(lambda x: jnp.dot(x, x.T))(mats)
result.shape

# compute the mean on each device in parallel and print the results
print(pmap(jnp.mean)(result))

from jax import lax

normalize = lambda x: x / lax.psum(x, axis_name='i')
result = pmap(normalize, axis_name='i')(jnp.arange(4.))
print(result)

from functools import partial

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

print(normalize(jnp.arange(4.)))

@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def f(x):
  row_normed = x / lax.psum(x, 'rows')
  col_normed = x / lax.psum(x, 'cols')
  doubly_normed = x / lax.psum(x, ('rows', 'cols'))
  return row_normed, col_normed, doubly_normed

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

print(a)
print(a.sum(0))

from jax._src import xla_bridge
device_count = jax.device_count()

def send_right(x, axis_name):
  left_perm = [(i, (i + 1) % device_count) for i in range(device_count)]
  return lax.ppermute(x, perm=left_perm, axis_name=axis_name)

def send_left(x, axis_name):
  left_perm = [((i + 1) % device_count, i) for i in range(device_count)]
  return lax.ppermute(x, perm=left_perm, axis_name=axis_name)

def update_board(board):
  left = board[:-2]
  right = board[2:]
  center = board[1:-1]
  return lax.bitwise_xor(left, lax.bitwise_or(center, right))

@partial(pmap, axis_name='i')
def step(board_slice):
  left, right = board_slice[:1], board_slice[-1:]
  right, left = send_left(left, 'i'), send_right(right, 'i')
  enlarged_board_slice = jnp.concatenate([left, board_slice, right])
  return update_board(enlarged_board_slice)

def print_board(board):
  print(''.join('*' if x else ' ' for x in board.ravel()))


board = np.zeros(40, dtype=bool)
board[board.shape[0] // 2] = True
reshaped_board = board.reshape((device_count, -1))

print_board(reshaped_board)
for _ in range(20):
  reshaped_board = step(reshaped_board)
  print_board(reshaped_board)

from jax import grad

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

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