Installing JAX

# CPU-only (works on all platforms)
pip install jax

# GPU support with CUDA 12
pip install jax[cuda12]

# Or specify CUDA version explicitly
pip install jax[cuda12_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# On a TPU VM (Google Cloud)
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# Neural network libraries
pip install flax       # Google's neural network library for JAX
pip install optax      # Gradient processing and optimization

# Alternative neural network library
pip install dm-haiku   # DeepMind's Haiku

# Useful extras
pip install orbax-checkpoint  # Model checkpointing
pip install clu               # Common loop utilities

import jax
import jax.numpy as jnp

# Check version
print(f"JAX version: {jax.__version__}")

# Check available devices
print(f"Devices: {jax.devices()}")

# Quick test: create an array and compute
x = jnp.array([1.0, 2.0, 3.0])
print(f"Sum: {jnp.sum(x)}")
print(f"Device: {x.devices()}")

# Test JIT compilation
@jax.jit
def f(x):
    return jnp.dot(x, x)

result = f(jnp.ones(1000))
print(f"JIT works! Result: {result}")