Chapter 1: Introduction to JAX

# Imports and environment info used throughout the chapter
import sys
import time
import warnings
import subprocess

import numpy as np
import jax
import jax.numpy as jnp

print(f"Python : {sys.version.split()[0]}")
print(f"JAX    : {jax.__version__}")
print(f"NumPy  : {np.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices: {jax.devices()}")

# Check: jax.shard_map is the currently recommended parallel API, and the jax.pmap docstring marks it as legacy
print(jax.pmap.__doc__.strip().splitlines()[0])
assert hasattr(jax, "shard_map")
assert "shard_map" in jax.pmap.__doc__

# Look at the jaxpr that JAX gets from tracing
def f(x):
    return jnp.sin(x) ** 2 + 1.0

print("jaxpr of f:")
print(jax.make_jaxpr(f)(1.0))

print("jaxpr of grad(f) (the extra cos and multiplications are the chain rule):")
print(jax.make_jaxpr(jax.grad(f))(1.0))

# Check: d/dx [sin²x + 1] = 2·sin x·cos x
assert jnp.allclose(jax.grad(f)(1.0), 2 * jnp.sin(1.0) * jnp.cos(1.0))

# Check available devices (the CPU is always there)
print(f"JAX devices: {jax.devices()}")

# Creating a NumPy array
np_array = np.array([1.0, 2.0, 3.0])
print(f"NumPy array: {np_array}, type: {type(np_array)}")

# Creating a JAX array: check its type with jax.Array (the concrete class is ArrayImpl)
jnp_array = jnp.array([1.0, 2.0, 3.0])
print(f"JAX array: {jnp_array}, type: {type(jnp_array)}")

assert isinstance(jnp_array, jax.Array)
assert not isinstance(jnp_array, np.ndarray)   # a JAX array is not a NumPy array

# Basic operations are written exactly the same way
a_np = np.arange(6).reshape((2, 3))
b_np = np.array([[1, 1, 1], [2, 2, 2]])
c_np = a_np + b_np * 2

a_jnp = jnp.arange(6).reshape((2, 3))
b_jnp = jnp.array([[1, 1, 1], [2, 2, 2]])
c_jnp = a_jnp + b_jnp * 2  # same syntax as NumPy

print(f"NumPy result:\n{c_np}")
print(f"JAX result:\n{c_jnp}")

# Check that the results are close (useful when comparing floating-point results)
print(f"Results close: {np.allclose(c_np, c_jnp)}")
assert np.allclose(c_np, c_jnp)

np_array = np.array([1, 2, 3])
np_array[0] = 100  # fine in NumPy
print(f"Modified NumPy array: {np_array}")

jnp_array = jnp.array([1, 2, 3])
try:
    # raises TypeError in JAX
    jnp_array[0] = 100
except TypeError as e:
    print(f"\nError when modifying a JAX array in place: {e}")

# The JAX way: .at[index].set(value) creates a new array with the new value
updated_jnp_array = jnp_array.at[0].set(100)
print(f"\nOriginal JAX array (unchanged): {jnp_array}")
print(f"Updated JAX array (new object): {updated_jnp_array}")

assert jnp.array_equal(jnp_array, jnp.array([1, 2, 3]))
assert jnp.array_equal(updated_jnp_array, jnp.array([100, 2, 3]))

# Asynchronous dispatch: not waiting vs. waiting for the computation to finish
x = jax.random.normal(jax.random.key(0), (2000, 2000))
_ = (x @ x).block_until_ready()  # warm up to exclude first-call overhead

t0 = time.perf_counter()
y = x @ x                        # only dispatches, returns immediately
t_dispatch = time.perf_counter() - t0
y.block_until_ready()            # wait for it to finish so it doesn't affect the next timing

t0 = time.perf_counter()
y = (x @ x).block_until_ready()  # wait until it is really done
t_real = time.perf_counter() - t0

print(f"No wait (dispatch only): {t_dispatch * 1e3:8.3f} ms")
print(f"block_until_ready: {t_real * 1e3:8.3f} ms")
print("Takeaway: benchmarks must use .block_until_ready(), or they badly underestimate the time.")
assert t_dispatch < t_real

# Check the default floating-point type
np_float = np.array([1.0]).dtype
jnp_float = jnp.array([1.0]).dtype
print(f"Default NumPy float type: {np_float}")
print(f"Default JAX float type: {jnp_float}")
assert np_float == np.float64
assert jnp_float == jnp.float32

# ❌ Old style (outdated): with the default config, dtype=jnp.float64 has no effect and is truncated to float32
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    try:
        jnp_float64 = jnp.array([1.0, 2.0], dtype=jnp.float64)
        print(f"Requested float64, actual dtype: {jnp_float64.dtype}")
    except Exception as e:
        print(f"Error: {type(e).__name__}: {e}")
for w in caught:
    print(f"Observed warning: {w.message}")

assert jnp_float64.dtype == jnp.float32
assert any("truncated to dtype float32" in str(w.message) for w in caught)

# ✅ Recommended: turn on jax_enable_x64 when the program starts
# This must be set before creating any arrays; to avoid affecting the rest of this notebook, we demo it in a subprocess
x64_demo = '''
import jax
jax.config.update("jax_enable_x64", True)   # must be set before creating arrays
import jax.numpy as jnp
a = jnp.array([1.0, 2.0], dtype=jnp.float64)
assert a.dtype == jnp.float64
assert jnp.array([1.0]).dtype == jnp.float64 and jnp.arange(3).dtype == jnp.int64
print("With x64 on, explicit float64:", a.dtype)
print("With x64 on, default float:", jnp.array([1.0]).dtype, ", default int:", jnp.arange(3).dtype)
'''
r = subprocess.run([sys.executable, "-c", x64_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "float64" in r.stdout

# For comparison: turning on x64 midway -- no error, but arrays created before it are still float32
late_demo = '''
import jax, jax.numpy as jnp
early = jnp.ones(2)                          # created before x64 is on
jax.config.update("jax_enable_x64", True)    # deliberately turned on one step late
late = jnp.ones(2, dtype=jnp.float64)
assert early.dtype == jnp.float32 and late.dtype == jnp.float64
print("x64 turned on midway: array created before", early.dtype, ", array created after", late.dtype)
'''
r = subprocess.run([sys.executable, "-c", late_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "float32" in r.stdout

# Weak types: Python scalars don't raise precision to 64 bits
print("NumPy: int64 array + 1.5 ->", (np.array([1, 2, 3]) + 1.5).dtype)
print("JAX  : int32 array + 1.5 ->", (jnp.array([1, 2, 3]) + 1.5).dtype)
assert (jnp.array([1, 2, 3]) + 1.5).dtype == jnp.float32

# Integer overflow: NumPy raises, JAX silently wraps around
try:
    print(np.array([1, 2, 3], dtype=np.int8) + 300)
except OverflowError as e:
    print(f"\nNumPy int8 + 300 -> OverflowError: {e}")
wrapped = jnp.array([1, 2, 3], dtype=jnp.int8) + 300
print("JAX   int8 + 300 ->", wrapped, "(300 is truncated to 44, so the result is wrong)")
assert jnp.array_equal(wrapped, jnp.array([45, 46, 47], dtype=jnp.int8))

a = jnp.arange(5)
try:
    np.arange(5)[10]
except IndexError as e:
    print(f"NumPy out of bounds: IndexError: {e}")

print(f"JAX a[10]            -> {a[10]}   (clamped to the last element)")
print(f"JAX a.at[10].set(99) -> {a.at[10].set(99)}   (update dropped)")
print(f"JAX a.at[10].get(mode='fill', fill_value=-1) -> {a.at[10].get(mode='fill', fill_value=-1)}")

assert a[10] == 4
assert jnp.array_equal(a.at[10].set(99), a)
assert a.at[10].get(mode="fill", fill_value=-1) == -1

src = np.array([1.0, 2.0])
j = jnp.asarray(src)
src[0] = 99.0
print(f"After modifying the NumPy source array, the JAX array is still: {j}")
assert j[0] == 1.0

back = np.asarray(jnp.ones(2))
print(f"Is np.asarray(jax array) writable: {back.flags.writeable}")
try:
    back[0] = 5
except ValueError as e:
    print(f"Write error: {e}")
assert not back.flags.writeable
assert np.array(jnp.ones(2)).flags.writeable   # np.array gives a writable copy

def predict(params, inputs):
    # A simple linear model
    return jnp.dot(inputs, params['w']) + params['b']

def loss_fn(params, inputs, targets):
    predictions = predict(params, inputs)
    error = predictions - targets
    return jnp.mean(error ** 2)  # mean squared error

# grad: differentiate loss_fn with respect to its first argument (params)
grad_fn = jax.grad(loss_fn)

# jit: compile the gradient function
jit_grad_fn = jax.jit(grad_fn)

# vmap: suppose predict handles a single input; vmap lets it handle a batch
# Map over axis 0 of the inputs argument; params is shared across all examples (None)
batched_predict = jax.vmap(predict, in_axes=(None, 0))

# Build a small dataset: y = 3*x0 - 2*x1 + 0.5 + noise
key = jax.random.key(0)
k_x, k_noise = jax.random.split(key)
inputs = jax.random.normal(k_x, (8, 2))                          # 8 examples, 2 features each
true_w, true_b = jnp.array([3.0, -2.0]), 0.5
targets = inputs @ true_w + true_b + 0.1 * jax.random.normal(k_noise, (8,))

params = {'w': jnp.zeros(2), 'b': jnp.array(0.0)}

print("Loss:", loss_fn(params, inputs, targets))

grads = grad_fn(params, inputs, targets)
print("Gradients (a dict with the same structure as params):", grads)

jit_grads = jit_grad_fn(params, inputs, targets)
print("JIT gradients:", jit_grads)

# Check 1: results match with and without jit
assert jax.tree.all(jax.tree.map(jnp.allclose, grads, jit_grads))

# Check 2: matches the analytic gradient
resid = predict(params, inputs) - targets
assert jnp.allclose(grads['w'], 2 / len(targets) * inputs.T @ resid, atol=1e-5)
assert jnp.allclose(grads['b'], 2 * resid.mean(), atol=1e-5)

# vmap version: the single-example predict applied to a batch of inputs
true_params = {'w': true_w, 'b': jnp.array(true_b)}
single = predict(true_params, inputs[0])          # one example -> scalar
batch = batched_predict(true_params, inputs)      # a batch of examples -> (8,)
print("Single-example output shape:", single.shape, "batch output shape:", batch.shape)

# Check 3: the vmap result matches a per-example loop
assert batch.shape == (8,)
assert jnp.allclose(batch, jnp.stack([predict(true_params, xi) for xi in inputs]))

@jax.jit
def update(params, inputs, targets, lr=0.1):
    grads = jax.grad(loss_fn)(params, inputs, targets)
    # jax.tree.map applies the same update to every entry in the params dict
    return jax.tree.map(lambda p, g: p - lr * g, params, grads)

p = {'w': jnp.zeros(2), 'b': jnp.array(0.0)}
for step in range(201):
    p = update(p, inputs, targets)
    if step % 50 == 0:
        print(f"step {step:3d}  loss = {loss_fn(p, inputs, targets):.5f}")
print("Learned parameters:", p, "\nTrue parameters: w = [3, -2], b = 0.5")

# Check: converges to the noise level and the parameters are close to the true values
assert loss_fn(p, inputs, targets) < 0.01
assert jnp.allclose(p['w'], true_w, atol=0.15)

counter = {"n": 0}

@jax.jit
def impure(x):
    print("  [Python print] tracing, x =", x)        # runs during tracing; x is a Tracer
    jax.debug.print("  [debug.print] runtime x = {}", x)  # runs on every call
    counter["n"] += 1                                    # side effect: happens only during tracing
    return x * 2

print("Call 1 (shape (2,))"); impure(jnp.ones(2)).block_until_ready()
print("Call 2 (shape (2,))"); impure(jnp.ones(2)).block_until_ready()
print("Call 3 (shape (3,))"); impure(jnp.ones(3)).block_until_ready()
print(f"counter = {counter['n']} (called 3 times, but traced only 2 times)")
assert counter["n"] == 2

# For comparison: "abstract values with only shape and dtype" is jit's behavior; with grad alone, tracers carry concrete values
def py_if(x):
    return x * x if x > 0 else -x

assert jax.grad(py_if)(-2.0) == -1.0      # grad alone: Python if works
try:
    jax.jit(py_if)(-2.0)
    jit_if_ok = True
except jax.errors.TracerBoolConversionError:
    jit_if_ok = False
assert not jit_if_ok                       # jit: an abstract value can't be converted to bool
print("Python if works under grad alone; under jit it raises TracerBoolConversionError")

scale = 2.0

@jax.jit
def scaled(x):
    return x * scale   # reads a global variable -- impure

print("With scale=2:", scaled(jnp.array(1.0)))
scale = 100.0
print("After changing scale to 100:", scaled(jnp.array(1.0)), " <- still 2; it uses the old value from tracing time")
assert scaled(jnp.array(1.0)) == 2.0

# The right way: pass values that can change in as arguments
scaled_ok = jax.jit(lambda x, s: x * s)
assert scaled_ok(jnp.array(1.0), scale) == 100.0

loss_val, grads2 = jax.value_and_grad(loss_fn)(params, inputs, targets)
print("Loss:", loss_val)
print("Gradients:", grads2)
assert jnp.allclose(loss_val, loss_fn(params, inputs, targets))

# Differentiating with respect to several arguments: argnums=(0, 1) covers both params and inputs
g_params, g_inputs = jax.grad(loss_fn, argnums=(0, 1))(params, inputs, targets)
print("Shape of the gradient w.r.t. inputs:", g_inputs.shape)
assert g_inputs.shape == inputs.shape

# A non-scalar output raises an error
try:
    jax.grad(predict)(params, inputs)
except TypeError as e:
    print("Error when differentiating a non-scalar output:", str(e).splitlines()[0])

# Check: read the Python version requirement and the available install extras from the installed jax package metadata
import importlib.metadata as md_

meta = md_.metadata("jax")
extras = meta.get_all("Provides-Extra")
print("jax version:", meta["Version"])
print("Requires-Python:", meta["Requires-Python"])
print("Available install extras (pip install 'jax[xxx]'):", extras)

assert meta["Requires-Python"] == ">=3.12"
assert {"cuda12", "cuda13", "tpu"} <= set(extras)
assert "cuda11_cudnn82" not in extras
assert "jaxlib" in " ".join(meta.get_all("Requires-Dist"))   # jaxlib is installed automatically as a dependency

# ❌ Old style (outdated): the old PRNGKey; also, the example output 10.811179161071777 can't be reproduced on JAX 0.12
COURSE_VALUE = 10.811179161071777
try:
    key_old = jax.random.PRNGKey(0)
    x_old = jax.random.normal(key_old, (10,))
    y_course = float(jnp.dot(x_old, x_old))
    print(f"Value given in the original example: {COURSE_VALUE}")
    print(f"Actual output                      : {y_course}")
except Exception as e:
    print(f"Error: {type(e).__name__}: {e}")

assert abs(y_course - COURSE_VALUE) > 1e-3   # the value in the example can't be reproduced

# ✅ Recommended: an installation check script (using jax.random.key)
try:
    devices = jax.devices()
    print("JAX found these devices:")
    for i, device in enumerate(devices):
        print(f"{i}: {device.platform.upper()} ({device.device_kind})")

    # Try a simple computation
    key = jax.random.key(0)
    x = jax.random.normal(key, (10,))
    y = jnp.dot(x, x)
    print(f"\nRan a simple JAX operation successfully. Result: {y}")

    # Check the default backend
    print(f"\nDefault backend: {jax.default_backend()}")

except Exception as e:
    print("An error occurred while verifying JAX:")
    print(e)
    print("\nPlease check the installation steps, especially the GPU driver / CUDA version (if applicable).")

# Check: the result matches an independent NumPy computation, and the old and new key styles give the same random numbers
assert np.isclose(float(y), np.dot(np.asarray(x), np.asarray(x)), rtol=1e-5)
assert np.isclose(float(y), y_course)

# Extra demo: the exact random values depend on the RNG configuration (toggle jax_threefry_partitionable in a subprocess)
rng_demo = '''
import jax, jax.numpy as jnp
jax.config.update("jax_threefry_partitionable", {flag})
x = jax.random.normal(jax.random.key(0), (10,))
print(float(jnp.dot(x, x)))
'''
values = {}
for flag in (True, False):
    r = subprocess.run([sys.executable, "-c", rng_demo.format(flag=flag)], capture_output=True, text=True)
    assert r.returncode == 0, r.stderr
    values[flag] = float(r.stdout)
    print(f"jax_threefry_partitionable={flag!s:5}: dot(x, x) = {values[flag]}")
print(f"Value given in the original example   : {COURSE_VALUE}")

assert values[True] != values[False]                                  # different config, different random numbers
assert all(abs(v - COURSE_VALUE) > 1e-3 for v in values.values())     # neither matches the value in the example

jax.print_environment_info()

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

# Create an array from a list
python_list = [1.0, 2.0, 3.0]
jax_array_from_list = jnp.array(python_list)
print(f"Array created from a list: {jax_array_from_list}")
print(f"Type: {type(jax_array_from_list)}")
assert isinstance(jax_array_from_list, jax.Array)

# Create a 2D array from a nested list
python_nested_list = [[1, 2], [3, 4]]
jax_2d_array = jnp.array(python_nested_list)
print(f"2D array:\n{jax_2d_array}")

# Array of zeros
zeros_array = jnp.zeros((2, 3))  # shape (2 rows, 3 columns)
print(f"Zeros array:\n{zeros_array}")

# Array of ones
ones_array = jnp.ones((3,), dtype=jnp.int32)  # shape (3,), integer type
print(f"Ones array: {ones_array}")
print(f"Ones array dtype: {ones_array.dtype}")

# Evenly spaced integer sequence
range_array = jnp.arange(5)  # like Python's range(5) -> [0, 1, 2, 3, 4]
print(f"arange array: {range_array}")

# Evenly spaced points in an interval
linspace_array = jnp.linspace(0, 1, 5)  # 5 points from 0 to 1 (endpoints included)
print(f"linspace array: {linspace_array}")

assert zeros_array.dtype == jnp.float32 and ones_array.dtype == jnp.int32

from jax import random

# Create the initial key (recommended style); the seed is usually an integer
key = random.key(0)
print(f"Initial key: {key.dtype}, underlying data {random.key_data(key)}")

# Generate normally distributed random numbers. The random function "uses" the key but doesn't modify it
normal_random_numbers = random.normal(key, shape=(2, 2))
print(f"Normal random numbers:\n{normal_random_numbers}")

# To generate more random numbers, split the key first
key, subkey = random.split(key)  # gives a new key and a subkey
uniform_random_numbers = random.uniform(subkey, shape=(3,))
print(f"\nKey after splitting: {random.key_data(key)}")
print(f"Subkey: {random.key_data(subkey)}")
print(f"Uniform random numbers: {uniform_random_numbers}")

# Calling again with the "same" original key gives exactly the same result -- purity and reproducibility
same_normal_numbers = random.normal(random.key(0), shape=(2, 2))
print(f"\nSame normal random numbers:\n{same_normal_numbers}")
assert jnp.allclose(normal_random_numbers, same_normal_numbers)

# key vs. PRNGKey demo (apart from the ❌ comparison cell in 1.4, this is the only place PRNGKey appears)
old_key = jax.random.PRNGKey(0)
new_key = jax.random.key(0)
print(f"PRNGKey(0): {old_key}, dtype={old_key.dtype}, shape={old_key.shape}")
print(f"key(0)    : dtype={new_key.dtype}, shape={new_key.shape}")
print(f"key_data(key(0)) = {jax.random.key_data(new_key)}")

# They generate the same random numbers
same = jnp.array_equal(random.normal(old_key, (3,)), random.normal(new_key, (3,)))
print("Same random numbers:", same)

# Shapes after batch splitting
print("split(PRNGKey, 3).shape =", random.split(old_key, 3).shape)
print("split(key, 3).shape     =", random.split(new_key, 3).shape)

# An old-style key can be mistaken for data in arithmetic; a new-style key raises an error
print("Old-style key + 1 =", old_key + 1, "(silently succeeds, which is dangerous)")
try:
    new_key + 1
except TypeError as e:
    print("New-style key + 1 raises:", str(e).splitlines()[0])

assert same
assert old_key.dtype == jnp.uint32 and old_key.shape == (2,)
assert new_key.shape == () and jax.dtypes.issubdtype(new_key.dtype, jax.dtypes.prng_key)
assert random.split(new_key, 3).shape == (3,)
assert bool(jax.random.wrap_key_data(old_key) == new_key)   # converting between old and new

# Both styles can choose the RNG algorithm, but only the typed key records it in the dtype
raw_rbg = jax.random.PRNGKey(0, impl="rbg")
typed_rbg = jax.random.key(0, impl="rbg")
print("PRNGKey(0, impl='rbg'):", raw_rbg.dtype, raw_rbg.shape, "| key(0, impl='rbg'):", typed_rbg.dtype)
assert raw_rbg.dtype == jnp.uint32 and raw_rbg.shape == (4,)
assert "rbg" in str(typed_rbg.dtype)
try:
    random.normal(raw_rbg, (2,))      # raw data carries no algorithm info: it's read as the default threefry, and the shape doesn't match
    raw_ok = True
except TypeError:
    raw_ok = False
assert not raw_ok
random.normal(typed_rbg, (2,))        # a typed key carries its algorithm, so it just works

key = jax.random.key(42)

# Wrong: the same key used twice
w1 = random.normal(key, (3,))
w2 = random.normal(key, (3,))
print("Reused key -> the two sets of \"random\" numbers are identical:", jnp.array_equal(w1, w2))
assert jnp.array_equal(w1, w2)

# Right: split off independent subkeys
key, k1, k2 = random.split(key, 3)
w1, w2 = random.normal(k1, (3,)), random.normal(k2, (3,))
print("Split subkeys -> different:", not jnp.array_equal(w1, w2))
assert not jnp.array_equal(w1, w2)

# Derive keys from the step number with fold_in inside a loop
base = jax.random.key(0)
for step in range(3):
    print(f"step {step}:", random.uniform(jax.random.fold_in(base, step), (2,)))

# Switching to a different distribution doesn't help: uniform and normal from the same key sort in exactly the same order (monotonically related)
u_same, z_same = random.uniform(key, (1000,)), random.normal(key, (1000,))
print("uniform and normal from the same key sort the same way:", bool(jnp.all(jnp.argsort(u_same) == jnp.argsort(z_same))))
assert jnp.all(jnp.argsort(u_same) == jnp.argsort(z_same))

x = jnp.arange(12).reshape((3, 4))
print(f"Array x:\n{x}")

print(f"Shape: {x.shape}")      # tuple of dimension sizes (rows, columns)
print(f"Dtype: {x.dtype}")  # element type, e.g. float32, int32
print(f"Number of dimensions: {x.ndim}")   # number of axes (2 for a matrix)
print(f"Total elements: {x.size}")   # total number of elements, 3 * 4 = 12

assert x.shape == (3, 4) and x.dtype == jnp.int32 and x.ndim == 2 and x.size == 12

# ❌ Old style (outdated): the code uses np without import numpy as np, so running it directly raises NameError
course_src = '''
numpy_array = np.array([1, 2, 3])
numpy_array[0] = 100
'''
try:
    exec(course_src, {})   # run in an empty namespace to simulate copying this code out and running it alone
    course_ok = True
except NameError as e:
    course_ok = False
    print(f"Observed: NameError: {e}")
assert not course_ok

# ✅ Recommended: import NumPy first (this tutorial already does import numpy as np at the top)
import numpy as np

numpy_array = np.array([1, 2, 3])
numpy_array[0] = 100  # fine in NumPy
print(f"Modified NumPy array: {numpy_array}")

jax_array = jnp.array([1, 2, 3])
try:
    # raises an error, because JAX arrays are immutable
    jax_array[0] = 100
except TypeError as e:
    print(f"\nError when trying to modify a JAX array in place: {e}")

# Use the functional "indexed update" syntax instead, which returns a new array
updated_jax_array = jax_array.at[0].set(100)
print(f"\nOriginal JAX array (unchanged): {jax_array}")
print(f"Updated JAX array (new object): {updated_jax_array}")

# Other kinds of updates are available too
incremented_array = jax_array.at[1].add(10)  # add 10 to the element at index 1
print(f"Incremented array: {incremented_array}")

assert jnp.array_equal(jax_array, jnp.array([1, 2, 3]))
assert jnp.array_equal(updated_jax_array, jnp.array([100, 2, 3]))
assert jnp.array_equal(incremented_array, jnp.array([1, 12, 3]))

x = jnp.zeros(5)
print("set     :", x.at[1:3].set(7.0))
print("multiply:", jnp.arange(5.).at[::2].multiply(10))
print("max     :", jnp.arange(5.).at[:].max(2.0))
print("apply   :", jnp.arange(5.).at[2].apply(jnp.sqrt))

# Repeated indices: JAX accumulates, NumPy doesn't
idx = jnp.array([0, 0, 2])
jax_acc = jnp.zeros(3).at[idx].add(1)
print("\nJAX  .at[[0,0,2]].add(1):", jax_acc)
a_np = np.zeros(3); a_np[[0, 0, 2]] += 1
print("NumPy a[[0,0,2]] += 1   :", a_np)
b_np = np.zeros(3); np.add.at(b_np, [0, 0, 2], 1)
print("NumPy np.add.at          :", b_np)
assert jnp.array_equal(jax_acc, jnp.array([2., 0., 1.]))
assert np.array_equal(a_np, [1., 0., 1.]) and np.array_equal(b_np, jax_acc)

# Inside jit, use jnp.where instead of boolean-mask assignment
relu = jax.jit(lambda v: jnp.where(v > 0, v, 0.0))
out = relu(jnp.array([-1.0, 2.0, -3.0, 4.0]))
print("\nReLU with jnp.where:", out)
assert jnp.array_equal(out, jnp.array([0., 2., 0., 4.]))

# In-place updates: check the compiled HLO for copy instructions
def hlo_has_copy(fn, *args):
    return " copy(" in fn.lower(*args).compile().as_text()

v0 = jnp.zeros(1000)
assert not hlo_has_copy(jax.jit(lambda v: (v * 2).at[0].set(1.0)), v0)             # updating an intermediate: in place
assert hlo_has_copy(jax.jit(lambda v: v.at[0].set(1.0)), v0)                        # updating an input argument: copies first
assert not hlo_has_copy(jax.jit(lambda v: v.at[0].set(1.0), donate_argnums=0), v0)  # after donating the input: in place
print("\nHLO check: no copy when updating an intermediate; a copy when updating an input argument; no copy with donate_argnums")

a = jnp.array([[1, 2], [3, 4]])
b = jnp.array([[5, 6], [7, 8]])

# Element-wise operations
print(f"a + b:\n{a + b}")
print(f"a * b:\n{a * b}")  # element-wise multiplication

# Matrix multiplication
print(f"Matrix product (jnp.dot):\n{jnp.dot(a, b)}")
print(f"Matrix product (@):\n{a @ b}")

# Universal functions
print(f"Sine of a:\n{jnp.sin(a)}")
print(f"Exponential of a:\n{jnp.exp(a)}")

assert jnp.array_equal(a @ b, jnp.array([[19, 22], [43, 50]]))
assert np.allclose(jnp.exp(a), np.exp(np.asarray(a)))

data = jnp.arange(10)
print(f"Original data: {data}")

# Get a single element
print(f"Element at index 3: {data[3]}")

# Get a slice
print(f"Slice from index 2 to 5: {data[2:5]}")

# Multi-dimensional indexing
matrix = jnp.arange(9).reshape((3, 3))
print(f"Matrix:\n{matrix}")
print(f"Element at row 1, column 2: {matrix[1, 2]}")
print(f"First two rows:\n{matrix[:2, :]}")
print(f"First column:\n{matrix[:, 0]}")

assert data[3] == 3 and matrix[1, 2] == 5
assert jnp.array_equal(matrix[:, 0], jnp.array([0, 3, 6]))

# ❌ Old style (outdated): x.device() is now an attribute; an except AttributeError fallback doesn't catch the new error
x = jnp.ones(3)
escaped = None
try:
    try:
        print(f"Device type of array x: {x.device().platform}")
    except AttributeError:
        print("Device info may need a specific context, or check jax.devices().")   # fallback branch, never runs
except TypeError as e:
    escaped = e
    print(f"The TypeError was not caught by except AttributeError: {e}")
assert isinstance(escaped, TypeError)

# ✅ Recommended
x = jnp.ones(3)
print(f"x.device           -> {x.device} (attribute, no parentheses)")
print(f"x.device.platform  -> {x.device.platform}")
print(f"x.devices()        -> {x.devices()} (method, returns a set of devices)")
print(f"Available devices: {jax.devices()}")

assert isinstance(x.device, jax.Device)
assert x.devices() == {x.device}
assert x.device.platform == jax.default_backend()

# List all available devices
devices = jax.devices()
print(f"Available devices: {devices}")

# Default backend
print(f"Default backend: {jax.default_backend()}")

# Details for each device
for d in devices:
    print(f"  id={d.id}, platform={d.platform}, kind={d.device_kind}")
assert all(d.platform == jax.default_backend() for d in devices)

x = jnp.arange(10.0)
print(f"Array x is on: {x.device}")
assert x.device == jax.devices()[0]

# ❌ Old style (outdated): Array.device changed from a method to an attribute; calling it with parentheses raises TypeError
random_array = jax.random.uniform(jax.random.key(42), shape=(2, 2))
try:
    print("Device:", random_array.device())
    old_api_works = True
except TypeError as e:
    old_api_works = False
    print(f"Observed: TypeError: {e}")
assert not old_api_works

# ✅ Recommended: three ways to look
print("random_array.device    ->", random_array.device)
print("random_array.devices() ->", random_array.devices())
print("random_array.sharding  ->", random_array.sharding)

assert random_array.devices() == {random_array.device}
assert random_array.sharding.device_set == random_array.devices()
assert random_array.sharding.is_fully_replicated   # single-device array: all the data is on this one device

devices = jax.devices()
x = jnp.arange(10.0)

# Put it on the first device
x_dev0 = jax.device_put(x, devices[0])
print(f"x_dev0 is on: {x_dev0.device}")

# Put it on the second device (needs at least two devices)
try:
    x_dev1 = jax.device_put(x, devices[1])
    print(f"x_dev1 is on: {x_dev1.device}")
except IndexError:
    print(f"Only {len(devices)} device(s), can't place on devices[1]")

# Put it on the CPU explicitly (useful when you have a GPU)
cpu_device = jax.devices("cpu")[0]
x_cpu = jax.device_put(x, cpu_device)
print(f"x_cpu is on: {x_cpu.device}")
assert x_cpu.device == cpu_device

# Without a device specified, new arrays go on the default device
y = jnp.ones(5)
print(f"y is on the default device: {y.device}")

# Asking for a backend that doesn't exist raises RuntimeError
try:
    jax.devices("gpu")
except RuntimeError as e:
    print(f"\njax.devices('gpu') raises: {e}")

multi_dev_demo = '''
import jax
jax.config.update("jax_num_cpu_devices", 4)   # simulate 4 CPU devices; must be set before using JAX
import jax.numpy as jnp

devs = jax.devices()
assert len(devs) == 4
print("Simulated devices:", devs)

x = jnp.arange(4.0)
a = jax.device_put(x, devs[1])        # committed to cpu:1
b = jax.device_put(x, devs[2])        # committed to cpu:2
print("a is on", a.device, ", b is on", b.device)

assert (a + 1).device == devs[1]
assert (a + jnp.ones(4)).device == devs[1]          # the uncommitted array is moved to cpu:1 automatically
print("a + jnp.ones(4) (the uncommitted array is moved over automatically) is on", (a + jnp.ones(4)).device)

raised = False
try:
    a + b
except ValueError as e:
    raised = True
    print("a + b raises ValueError:", str(e)[:70], "...")
assert raised

# The right way: move the data onto the same device explicitly first
c = a + jax.device_put(b, devs[1])
assert c.device == devs[1] and jnp.array_equal(c, 2 * x)
print("After moving to the same device, a + b =", c, ", on", c.device)
print("ALL_OK")
'''
r = subprocess.run([sys.executable, "-c", multi_dev_demo], capture_output=True, text=True)
print(r.stdout or r.stderr)
assert r.returncode == 0 and "ALL_OK" in r.stdout

x = jnp.arange(4.0) * 2

host_np = np.asarray(x)
host_get = jax.device_get(x)
print("np.asarray     ->", type(host_np), host_np)
print("jax.device_get ->", type(host_get), host_get)
print("scalar .item() ->", x.sum().item(), type(x.sum().item()))
assert isinstance(host_get, np.ndarray) and isinstance(x.sum().item(), float)

with jax.default_device(jax.devices("cpu")[0]):
    z = jnp.zeros(2)
print("Array created inside the default_device block is on:", z.device)
assert z.device == jax.devices("cpu")[0]

# Combined exercise: array creation, attributes, and devices
# numpy is usually for comparison, or for operations JAX doesn't support (already imported at the top)
print(f"JAX version: {jax.__version__}")
print(f"JAX default backend: {jax.default_backend()}")
# Check available devices (depends on your environment)
print(f"Available JAX devices: {jax.devices()}")

# Create from a Python list
py_list = [1.0, 2.5, 3.0, 4.2]
jax_array_from_list = jnp.array(py_list)
print("JAX array created from a list:", jax_array_from_list)
print("Type:", type(jax_array_from_list))
print("Dtype:", jax_array_from_list.dtype)

# Check: the type is jax.Array, and the default on CPU is float32 too
assert isinstance(jax_array_from_list, jax.Array)
assert "DeviceArray" not in type(jax_array_from_list).__name__
assert jax_array_from_list.dtype == jnp.float32

# Create specific arrays
zeros_array = jnp.zeros((2, 3))  # shape (2 rows, 3 columns)
print("\nZeros array:\n", zeros_array)

ones_array = jnp.ones((3, 2), dtype=jnp.int32)  # specify the dtype
print("\nOnes array (int32):\n", ones_array)
print("Dtype:", ones_array.dtype)

range_array = jnp.arange(0, 10, 2)  # start, stop (exclusive), step
print("\nRange array:", range_array)

linspace_array = jnp.linspace(0, 1, 5)  # start, stop (inclusive), number of points
print("\nEvenly spaced array:", linspace_array)

assert jnp.array_equal(range_array, jnp.array([0, 2, 4, 6, 8]))
assert jnp.allclose(linspace_array, jnp.array([0, .25, .5, .75, 1.]))

# Create a PRNG key
key = jax.random.key(42)  # seed 42

# Generate uniform random numbers in [0, 1)
random_array = jax.random.uniform(key, shape=(2, 2))
print("Random array (uniform):\n", random_array)

# Important: to get new random numbers, you must "split" the key first
key, subkey = jax.random.split(key)
random_normal_array = jax.random.normal(subkey, shape=(3,))  # standard normal distribution
print("\nRandom array (normal):\n", random_normal_array)

assert ((random_array >= 0) & (random_array < 1)).all()

print("Inspecting random_array:")
print("Shape:", random_array.shape)
print("Size:", random_array.size)   # total number of elements
print("Number of dimensions:", random_array.ndim)
print("Dtype:", random_array.dtype)

# Check which device the array is on (attribute, no parentheses; see 1.6 for the old .device() style)
print("Device:", random_array.device)

assert random_array.shape == (2, 2) and random_array.size == 4 and random_array.ndim == 2
assert random_array.dtype == jnp.float32

a = jnp.array([1, 2, 3])
b = jnp.array([4, 5, 6])

# Element-wise operations
c_add = a + b
print("a + b =", c_add)

c_mul = a * b
print("a * b =", c_mul)

# Operations with scalars (broadcasting)
c_scalar_add = a + 10
print("a + 10 =", c_scalar_add)

c_scalar_mul = a * 2
print("a * 2 =", c_scalar_mul)

# Check that a is unchanged (immutability)
print("Original 'a':", a)
assert jnp.array_equal(a, jnp.array([1, 2, 3]))
assert jnp.array_equal(c_mul, jnp.array([4, 10, 18]))

# Matrix multiplication
mat_a = jnp.array([[1, 2], [3, 4]])
mat_b = jnp.array([[5, 6], [7, 8]])

mat_prod = mat_a @ mat_b               # same as jnp.matmul(mat_a, mat_b)
print("Matrix product mat_a @ mat_b:\n", mat_prod)
assert jnp.array_equal(mat_prod, jnp.matmul(mat_a, mat_b))
assert np.array_equal(mat_prod, np.asarray(mat_a) @ np.asarray(mat_b))

data = jnp.arange(12).reshape((3, 4))
print("Original data:\n", data)

# A single element
print("\nElement at row 1, column 2:", data[1, 2])

# A whole row
print("Row 0:", data[0, :])

# A whole column
print("Column 1:", data[:, 1])

# A sub-array
print("Sub-array (first 2 rows, columns 1-2):\n", data[:2, 1:3])

# Direct assignment isn't supported
try:
    data[0, 0] = 99
except TypeError as e:
    print("\nDirect assignment raises:", str(e).splitlines()[0])

# Functional update: returns a new array; this is the style that works with transformations like jit and grad
updated_data = data.at[0, 0].set(99)
print("\nUpdated data:\n", updated_data)
print("Original data (unchanged):\n", data)

assert data[0, 0] == 0 and updated_data[0, 0] == 99
assert jnp.array_equal(data[:2, 1:3], jnp.array([[1, 2], [5, 6]]))