import jax
import jaxlib

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

"""
Explanation: <a href="https://colab.research.google.com/github/google/jax/blob/main/tests/notebooks/colab_cpu.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
JAX Colab CPU Test
This notebook is meant to be run in a Colab CPU runtime as a basic check for JAX updates.
End of explanation
"""


from jaxlib import xla_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 device.platform == "cpu", f"unexpected JAX device type: {device.platform}"

"""
Explanation: Confirm Device
End of explanation
"""


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)

"""
Explanation: Matrix Multiplication
End of explanation
"""


import jax.numpy as jnp
import jax.random as rand

N = 10
M = 20
key = rand.PRNGKey(1701)

X = rand.normal(key, (N, M))
u, s, vt = jnp.linalg.svd(X)
assert u.shape == (N, N)
assert vt.shape == (M, M)
print(s)

"""
Explanation: Linear Algebra
End of explanation
"""


@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)

"""
Explanation: XLA Compilation
End of explanation
"""


