JAX Colab TPU Test¶

import jax
import jaxlib

!cat /var/colab/hostname
print(jax.__version__)
print(jaxlib.__version__)

import requests
import os
if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20200416'
  resp = requests.post(url)
  assert resp.status_code == 200
  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)

from jaxlib import tpu_client_extension
import jax
key = jax.random.PRNGKey(1701)
arr = jax.random.normal(key, (1000,))
device = arr.device_buffer.device()
print(f"JAX device type: {device}")
assert isinstance(device, tpu_client_extension.TpuDevice), "unexpected JAX device type"

import jax
import numpy as np

# matrix multiplication on GPU
key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (3000, 3000))
result = jax.numpy.dot(x, x.T).mean()
print(result)

@jax.jit
def selu(x, alpha=1.67, lmbda=1.05):
  return lmbda * jax.numpy.where(x > 0, x, alpha * jax.numpy.exp(x) - alpha)
x = jax.random.normal(key, (5000,))
result = selu(x).block_until_ready()
print(result)