Chapter 6: Managing State in JAX

import sys
import warnings
from functools import partial
from typing import NamedTuple
from dataclasses import dataclass, field

import jax
import jax.numpy as jnp
from jax import lax

print("Python", sys.version.split()[0])
print("JAX   ", jax.__version__)
print("Devices", jax.devices())

# Pure function: the output depends only on the inputs, and nothing outside is affected
def pure_add(a, b):
    return a + b

# Impure: printing is an I/O side effect
def impure_add_and_print(a, b):
    result = a + b
    print(f"The result is {result}")   # side effect
    return result

# Impure: modifies a global variable
call_count = 0
def impure_increment_counter(x):
    global call_count
    call_count += 1             # side effect: modifies global state
    return x + 1

# Impure: modifies an input argument in place
def impure_append_to_list(data_list, value):
    data_list.append(value)     # side effect: modifies the caller's list
    return data_list

print(pure_add(2, 3))
impure_add_and_print(2, 3)
impure_increment_counter(1); impure_increment_counter(1)
print("call_count =", call_count)
lst = [1, 2]
impure_append_to_list(lst, 3)
print("The caller's list was changed:", lst)

# A typical object-oriented stateful style: the state hides inside the object, and count() modifies self.n every time
class Counter:
    def __init__(self):
        self.n = 0

    def count(self):
        self.n += 1
        return self.n

    def reset(self):
        self.n = 0

my_counter = Counter()
print(my_counter.count())   # 1
print(my_counter.count())   # 2  (same call, different result: not a pure function)

counter = 0

@jax.jit
def double(x):
    global counter
    counter += 1          # side effect: runs only during tracing
    return x * 2

for i in range(3):
    double(jnp.float32(i))
print("After 3 calls, counter =", counter, "(incremented only during the first trace)")
double(jnp.ones(3))       # new shape → retrace
print("After a call with a different shape, counter =", counter)
assert counter == 2

# Reading a global variable: the value is frozen at trace time
scale = 2.0
@jax.jit
def scaled(x):
    return x * scale

print("With scale=2:", scaled(1.0))
scale = 10.0
print("After changing to scale=10, scalar input (cache hit):", scaled(1.0))
print("After changing to scale=10, new-shape input (recompiled):", scaled(jnp.ones(2)))
assert float(scaled(1.0)) == 2.0 and float(scaled(jnp.ones(2))[0]) == 10.0   # same function, old and new values side by side

# The right way to print for debugging
@jax.jit
def show(x):
    print("Python print: appears only during tracing, x =", x)
    jax.debug.print("jax.debug.print: appears on every run, x = {}", x)
    return x + 1

show(1.0)
_ = show(2.0)

# Standard Python: modifying a dict in place
my_data = {'values': [10, 20]}

def update_data(data_dict, new_value):
    data_dict['values'].append(new_value)   # modify the list in place
    data_dict['last_added'] = new_value     # add a key in place
    # no return needed: the caller's object has already been changed

update_data(my_data, 30)
print(my_data)   # {'values': [10, 20, 30], 'last_added': 30}

# ❌ Wrong: jit a function that modifies an object's attribute in place
class StateHolder:
    def __init__(self, value):
        self.value = jnp.array(value)

def update_state_impurely(state_obj, increment):
    # Impure: modifies the object's attribute in place
    state_obj.value = state_obj.value + increment
    return state_obj.value

# Without JIT: runs as ordinary Python and looks fine
my_state = StateHolder(10.0)
print("Non-JIT call 1:", update_state_impurely(my_state, 5.0))   # 15.0
print("Non-JIT call 2:", update_state_impurely(my_state, 3.0))   # 18.0

# With JIT: always raises TypeError
my_state_jit = StateHolder(10.0)
jitted_update = jax.jit(update_state_impurely)
stateholder_error = None
try:
    print("JIT call 1:", jitted_update(my_state_jit, 5.0), my_state_jit.value)
except Exception as e:
    stateholder_error = e
    print(f"Actual result → {type(e).__name__}: {str(e)[:200]}")

assert isinstance(stateholder_error, TypeError), "StateHolder should reliably raise TypeError under jit"

# ✅ Fix: explicit state passing. The state is a plain array, the function returns the new state, and no object is modified
def update_state_pure(value, increment):
    return value + increment

jitted_pure = jax.jit(update_state_pure)
v = jnp.array(10.0)
v = jitted_pure(v, 5.0); print("JIT call 1:", v)
v = jitted_pure(v, 3.0); print("JIT call 2:", v)
assert float(v) == 18.0

# Demo: once registered as a PyTree, in-place modification "fails silently"
@jax.tree_util.register_pytree_node_class
class StateHolderPytree:
    def __init__(self, value):
        self.value = value
    def tree_flatten(self):
        return (self.value,), None          # (children, auxiliary data)
    @classmethod
    def tree_unflatten(cls, aux, children):
        return cls(*children)

s = StateHolderPytree(jnp.array(10.0))
f = jax.jit(update_state_impurely)
r1 = f(s, 5.0); print("Return value:", r1, "  outside s.value:", s.value)
r2 = f(s, 3.0); print("Return value:", r2, "  outside s.value:", s.value, " ← not accumulated! A silent error")
assert float(s.value) == 10.0 and float(r2) == 13.0     # the outside object was never modified

# JAX arrays are immutable
x = jnp.zeros(3)
try:
    x[0] = 1.0
except TypeError as e:
    print("TypeError:", str(e)[:110], "...")
print("Functional update x.at[0].set(1.0) →", x.at[0].set(1.0), "; the original array is still", x)

# Side effects under vmap: run only once during tracing and get a BatchTracer standing for the whole batch; there are no "multiple instances racing"
side_effect_log = []
def append_and_double(v):
    side_effect_log.append(type(v).__name__)   # side effect: modifies an external list
    return v * 2
jax.vmap(append_and_double)(jnp.arange(5.0))
print("After vmap processed 5 examples, the external list was appended to", len(side_effect_log), "time(s):", side_effect_log)
assert side_effect_log == ["BatchTracer"]

# Imperative: the state hides inside an object
class Counter:
    def __init__(self):
        self.count = 0
    def increment(self):
        self.count += 1
        return self.count

c = Counter()
print(c.increment())   # 1
print(c.increment())   # 2
print("Final count:", c.count)   # 2

# Functional: the state is passed in and returned explicitly
def init_counter_state():
    """Initialize the state."""
    return 0

def increment(current_state):
    """Take the current state and return the new state."""
    print(f"Running increment function with state: {current_state}")  # for demonstration
    new_state = current_state + 1
    return new_state

state = init_counter_state()
print("Initial state:", state)
state = increment(state)
print("State after first increment:", state)
state = increment(state)
print("State after second increment:", state)

jitted_increment = jax.jit(increment)

state = init_counter_state()
# The first call triggers tracing and compilation, which runs the Python code (including the print)
state = jitted_increment(state)
print("State after JIT call 1:", state)
# Later calls use the compiled version, so the print inside no longer runs
state = jitted_increment(state)
print("State after JIT call 2:", state)
state = jitted_increment(state)
print("State after JIT call 3:", state)

params = {
    'layer1': {'weights': jnp.ones((3, 2)), 'bias': jnp.zeros((2,))},
    'layer2': {'weights': jnp.ones((2, 1)), 'bias': jnp.zeros((1,))},
}
print(jax.tree.map(lambda x: x.shape, params))

def update_params(params, grads, learning_rate):
    """Apply one SGD update to every leaf of the PyTree."""
    def sgd_update(param, grad):
        return param - learning_rate * grad
    return jax.tree.map(sgd_update, params, grads)

# Placeholder gradients (same structure as params)
grads = {
    'layer1': {'weights': jnp.full((3, 2), 0.1), 'bias': jnp.full((2,), 0.01)},
    'layer2': {'weights': jnp.full((2, 1), 0.2), 'bias': jnp.full((1,), 0.02)},
}
learning_rate = 0.01

updated_params = update_params(params, grads, learning_rate)
print("layer1 weights:\n", updated_params['layer1']['weights'])
print("layer2 bias:", updated_params['layer2']['bias'])

# jit update_params directly; structure and values match the uncompiled version
jit_update_params = jax.jit(update_params)
updated_jit = jit_update_params(params, grads, learning_rate)
same = jax.tree.all(jax.tree.map(jnp.allclose, updated_params, updated_jit))
print("jit result matches the non-jit result:", bool(same))
assert same

# Self-verifying the migration table: none of the top-level old aliases exist, and both the tree_util and jax.tree versions do
removed_top_level = ["tree_map", "tree_leaves", "tree_structure", "tree_flatten",
                     "tree_unflatten", "tree_transpose", "tree_reduce"]
for name in removed_top_level:
    assert not hasattr(jax, name), f"jax.{name} should have been removed"
assert not hasattr(jax.tree_util, "tree_multimap")

still_in_tree_util = ["tree_map", "tree_leaves", "tree_structure", "tree_flatten", "tree_unflatten",
                      "tree_all", "tree_reduce", "tree_transpose", "tree_map_with_path",
                      "tree_leaves_with_path", "tree_broadcast"]
for name in still_in_tree_util:
    assert hasattr(jax.tree_util, name), f"jax.tree_util.{name} should still be available"

new_api = ["map", "leaves", "structure", "flatten", "unflatten", "all", "reduce",
           "transpose", "map_with_path", "leaves_with_path", "broadcast"]
for name in new_api:
    assert hasattr(jax.tree, name), f"jax.tree.{name} should exist"
print("Migration table verified; the jax.tree module provides:", [n for n in dir(jax.tree) if not n.startswith("_")])

# ❌ Old style (outdated): the top-level alias jax.tree_map was removed in 0.12
try:
    jax.tree_map(lambda x: x * 2, {'a': 1.0})
except AttributeError as e:
    print("Actual result → AttributeError:", e)

# ✅ Recommended: jax.tree.*, one function at a time
import operator
tree = {'b': jnp.arange(3.), 'a': (1.0, None, [2.0, 3.0])}

# (1) map: one tree, or several trees (several trees must have the same structure)
print("map, one tree:", jax.tree.map(lambda x: x * 2, {'w': 1.0, 'b': (2.0, 3.0)}))
summed = jax.tree.map(lambda a, b, c: a + b + c, {'w': 1.0, 'b': 2.0}, {'w': 10.0, 'b': 20.0}, {'w': 100.0, 'b': 200.0})
print("map, three trees:", summed)
assert summed == {'w': 111.0, 'b': 222.0}
try:
    jax.tree.map(lambda a, b: a + b, {'w': 1.0}, {'w': 1.0, 'b': 2.0})
except ValueError as e:
    print("Mismatched structures → ValueError:", str(e).splitlines()[0])

# (2) leaves: None is an "empty subtree", not a leaf; dicts are flattened in sorted key order ('a' before 'b')
lv = jax.tree.leaves(tree)
print("leaves:", lv)
assert len(lv) == 4 and lv[0] == 1.0

# (3) structure
print("structure:", jax.tree.structure(tree))

# (4) flatten / unflatten: a round trip
leaves, treedef = jax.tree.flatten(tree)
rebuilt = jax.tree.unflatten(treedef, [x * 10 for x in leaves])
print("unflatten:", rebuilt)
assert jax.tree.structure(rebuilt) == treedef

# (5) map_with_path: get the key path of each leaf, e.g. to apply weight decay to weights only
decayed = jax.tree.map_with_path(
    lambda path, x: x * 0.9 if 'weights' in jax.tree_util.keystr(path) else x, params)
print("map_with_path paths:", jax.tree.map_with_path(lambda p, x: jax.tree_util.keystr(p), params))
assert jnp.allclose(decayed['layer1']['weights'], 0.9) and jnp.allclose(decayed['layer1']['bias'], 0.0)

# (6) reduce: reduce all leaves to a single value, e.g. count parameters or compute the global gradient norm
n_params = jax.tree.reduce(operator.add, jax.tree.map(lambda x: x.size, params))
global_norm = jnp.sqrt(jax.tree.reduce(operator.add, jax.tree.map(lambda g: jnp.sum(g ** 2), grads)))
print("reduce, parameter count:", n_params, " global gradient norm:", global_norm)
assert n_params == 11

# (7) all: true only if every leaf is true
print("all:", jax.tree.all({'x': True, 'y': [True, True]}), jax.tree.all({'x': True, 'y': [True, False]}))
assert jax.tree.all({'x': True, 'y': [True, True]}) and not jax.tree.all({'x': True, 'y': [True, False]})

# (8) transpose: "list of dicts" ↔ "dict of lists", handy for regrouping per-step metrics by name
records = [{'loss': 1.0, 'acc': 0.5}, {'loss': 0.8, 'acc': 0.6}, {'loss': 0.5, 'acc': 0.7}]
outer = jax.tree.structure([0] * len(records))
inner = jax.tree.structure({'loss': 0, 'acc': 0})
by_key = jax.tree.transpose(outer, inner, records)
print("transpose:", by_key)
assert by_key == {'acc': [0.5, 0.6, 0.7], 'loss': [1.0, 0.8, 0.5]}

# 1) NamedTuple
class TrainState(NamedTuple):
    step: jax.Array
    params: dict

ts = TrainState(step=jnp.array(0), params={'w': jnp.ones(2)})
ts2 = ts._replace(step=ts.step + 1)        # creates a new object; the original is unchanged
print("NamedTuple structure:", jax.tree.structure(ts))
print("ts.step =", ts.step, " ts2.step =", ts2.step)

# 2) register_dataclass: specify the data fields and metadata fields explicitly
@partial(jax.tree_util.register_dataclass, data_fields=['w', 'b'], meta_fields=['name'])
@dataclass(frozen=True)
class Linear:
    w: jax.Array
    b: jax.Array
    name: str          # static metadata, not a leaf

lin = Linear(jnp.ones(2), jnp.zeros(2), 'fc1')
print("dataclass leaves:", jax.tree.leaves(lin))
print("dataclass structure:", jax.tree.structure(lin))

# You can also use it directly as a decorator and let the field metadata decide what's static
@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class AdamConfig:
    lr: jax.Array
    b1: float = field(default=0.9, metadata=dict(static=True))

print("Inferred:", jax.tree.structure(AdamConfig(jnp.array(1e-3))))

# frozen=True: an attempt to modify in place errors immediately
try:
    lin.w = jnp.zeros(2)
except Exception as e:
    print(f"{type(e).__name__}: {e}")

# All of these containers pass straight through jit / grad
@jax.jit
def forward(layer, x):
    return x @ layer.w + layer.b.sum()

print("jit forward:", forward(lin, jnp.array([1.0, 2.0])))
print("Gradient with respect to the dataclass:", jax.grad(lambda l: forward(l, jnp.array([1.0, 2.0])))(lin))

class Counter:
    def __init__(self):
        self.count = 0
    def increment(self):
        self.count += 1     # in-place modification: a side effect
        return self.count

counter_obj = Counter()
print(counter_obj.increment(), counter_obj.increment())

initial_state = 0

def update_counter(current_state):
    """Increment the input state by 1."""
    print("Tracing update_counter...")   # shows up only during tracing (for demonstration)
    return current_state + 1

state1 = update_counter(initial_state)
state2 = update_counter(state1)
state3 = update_counter(state2)
print("Final state:", state3)

counter_traces = []
def update_counter_counted(current_state):
    counter_traces.append(1)          # runs only during tracing; used to count traces
    return update_counter(current_state)

jitted_update_counter = jax.jit(update_counter_counted)

s = initial_state
for i in range(3):
    s = jitted_update_counter(s)
    print(f"After call {i + 1}: {s}")
assert len(counter_traces) == 1 and int(s) == 3

def scan_body(carry_state, _):
    next_state = update_counter(carry_state)
    return next_state, None          # no extra output per step

initial_scan_state = 0
num_steps = 10
final_state, accumulated_outputs = jax.lax.scan(
    scan_body, initial_scan_state, xs=None, length=num_steps)
print("Final state after scan:", final_state)
print("Per-step outputs:", accumulated_outputs)   # each step outputs None, so the whole thing is None too

# ❌ Wrong: steps is passed into jit as an ordinary argument and then into scan's length
@jax.jit
def run_scan_jitted(init_state, steps):
    final, _ = jax.lax.scan(scan_body, init_state, xs=None, length=steps)
    return final

scan_error = None
try:
    print(run_scan_jitted(initial_scan_state, num_steps))
except Exception as e:
    scan_error = e
    print(f"Actual result → {type(e).__name__}:\n{str(e)[:250]}")

assert isinstance(scan_error, jax.errors.ConcretizationTypeError)

# ✅ Fix 1: mark steps as static with static_argnames (each distinct value of steps compiles once)
@partial(jax.jit, static_argnames="steps")
def run_scan_jitted(init_state, steps):
    final, _ = jax.lax.scan(scan_body, init_state, xs=None, length=steps)
    return final

r_static = run_scan_jitted(initial_scan_state, steps=num_steps)
print("Fix 1 (static_argnames):", r_static)
assert int(r_static) == 10

# ✅ Fix 2: when the loop count needs to change at run time and you don't need per-step outputs, use fori_loop
@jax.jit
def run_fori(init_state, steps):
    return lax.fori_loop(0, steps, lambda i, c: update_counter(c), init_state)

r_fori_10, r_fori_25 = run_fori(initial_scan_state, 10), run_fori(initial_scan_state, 25)
print("Fix 2 (fori_loop):", r_fori_10, r_fori_25)
assert int(r_fori_10) == 10 and int(r_fori_25) == 25

def running_mean_body(carry, x):
    new_carry = {'sum': carry['sum'] + x, 'n': carry['n'] + 1}
    y = new_carry['sum'] / new_carry['n']        # per-step output: the current mean
    return new_carry, y

xs = jnp.array([2.0, 4.0, 6.0, 8.0, 10.0])
init = {'sum': jnp.array(0.0), 'n': jnp.array(0)}
final, means = lax.scan(running_mean_body, init, xs)
print("Final state:", final)
print("Mean at each step:", means)
assert jnp.allclose(means, jnp.array([2.0, 3.0, 4.0, 5.0, 6.0]))

# A carry whose dtype changes raises an error
def add_body(carry, x):
    return carry + x, None          # int state + float input → float

try:
    lax.scan(add_body, jnp.array(0, dtype=jnp.int32), xs)    # explicit int32 initial state
except TypeError as e:
    print("dtype mismatch → TypeError:", next(l for l in str(e).splitlines() if "input carry" in l)[:150])

# Exception: weak types (Python scalars, values like jnp.array(0) with no dtype given) are allowed to be promoted
print("Weakly typed initial value jnp.array(0):", lax.scan(add_body, jnp.array(0), xs)[0], "(no error)")

# A shape mismatch raises an error too
try:
    lax.scan(lambda c, x: (jnp.stack([c, c]), None), jnp.array(0.0), xs)
except TypeError as e:
    print("shape mismatch → TypeError:", next(l for l in str(e).splitlines() if "input carry" in l)[:150])

def objective_function(x):
    return (x - 3.0) ** 2

grad_fn = jax.grad(objective_function)

def sgd_update(params, gradients, learning_rate):
    """Pure function: return new parameters based on the gradients."""
    return params - learning_rate * gradients

current_params = jnp.array(0.0)
learning_rate = 0.1
num_steps = 20

print(f"Initial parameter: {current_params:.4f}")
for step in range(num_steps):
    gradients = grad_fn(current_params)
    next_params = sgd_update(current_params, gradients, learning_rate)
    current_params = next_params          # explicitly catch the new state
    if (step + 1) % 5 == 0:
        print(f"Step {step + 1:2d}: parameter = {current_params:.4f}, gradient = {gradients:.4f}")
print(f"Final optimized parameter: {current_params:.4f}")
sgd_final = float(current_params)       # keep it to compare with momentum later

# ❌ Poor hyperparameter: momentum=0.9 is underdamped on this problem and doesn't converge within 20 steps
class OptimizerState(NamedTuple):
    params: jax.Array
    velocity: jax.Array

def momentum_update(state, gradients, learning_rate, momentum):
    """Pure function: take the old optimizer state and return the new optimizer state."""
    new_velocity = momentum * state.velocity + learning_rate * gradients
    new_params = state.params - new_velocity
    return OptimizerState(params=new_params, velocity=new_velocity)

initial_params = jnp.array(0.0)
initial_velocity = jnp.array(0.0)
current_state = OptimizerState(params=initial_params, velocity=initial_velocity)

learning_rate = 0.1
momentum_coeff = 0.9
num_steps = 20

print(f"Initial state: parameter = {current_state.params:.4f}, velocity = {current_state.velocity:.4f}")
for step in range(num_steps):
    gradients = grad_fn(current_state.params)
    current_state = momentum_update(current_state, gradients, learning_rate, momentum_coeff)
    if (step + 1) % 5 == 0:
        print(f"Step {step + 1:2d}: parameter = {current_state.params:.4f}, "
              f"velocity = {current_state.velocity:.4f}, gradient = {gradients:.4f}")
print(f"Final optimized parameter: {current_state.params:.4f}")
momentum09_final = float(current_state.params)
print(f"Measured: |momentum μ=0.9 − 3| = {abs(momentum09_final - 3):.4f}  >  |plain SGD − 3| = {abs(sgd_final - 3):.4f}")
assert abs(momentum09_final - 3) > abs(sgd_final - 3)     # momentum with μ=0.9 actually does worse than plain SGD

# ✅ Fix: the same momentum_update, with only the momentum coefficient changed to 0.5 (learning rate unchanged)
def run_momentum(mu, lr=0.1, steps=20):
    st = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
    for _ in range(steps):
        st = momentum_update(st, grad_fn(st.params), lr, mu)
    return float(st.params)

for mu_ in [0.0, 0.3, 0.5, 0.9]:
    print(f"momentum={mu_}: parameter after 20 steps = {run_momentum(mu_):.4f}")

momentum05_final = run_momentum(0.5)
assert abs(momentum05_final - 3) < 1e-2                    # converges after tuning
assert abs(momentum05_final - 3) < abs(sgd_final - 3)      # and beats plain SGD

# ❌ Wrong: learning rate and momentum as static arguments (recompiles every step with a learning rate schedule)
momentum_traces = []
def momentum_update_logged(state, gradients, learning_rate, momentum):
    momentum_traces.append(learning_rate)        # runs only during tracing; used to count traces
    return momentum_update(state, gradients, learning_rate, momentum)

jitted_momentum_update = jax.jit(momentum_update_logged, static_argnums=(2, 3))

schedule = [0.1, 0.09, 0.08, 0.07, 0.06]   # simulated learning rate decay
current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
print("--- Running the JIT-compiled update (static learning rate) ---")
for lr_t in schedule:
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), lr_t, 0.5)
print(f"parameter = {current_state.params:.4f}; traces = {len(momentum_traces)}")
assert len(momentum_traces) == 5

# ✅ Fix: pass the learning rate and momentum as ordinary (dynamic) arguments
momentum_traces.clear()
jitted_momentum_update = jax.jit(momentum_update_logged)

current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
print("--- Running the JIT-compiled update (dynamic learning rate) ---")
for lr_t in schedule:
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), lr_t, 0.5)
print(f"parameter = {current_state.params:.4f}; traces = {len(momentum_traces)}")
assert len(momentum_traces) == 1

# Constant learning rate, full 20 steps: matches the uncompiled version
current_state = OptimizerState(params=jnp.array(0.0), velocity=jnp.array(0.0))
for step in range(num_steps):
    current_state = jitted_momentum_update(current_state, grad_fn(current_state.params), 0.1, 0.5)
print(f"Final optimized parameter (JIT-compiled, μ=0.5): {current_state.params:.4f}")
assert jnp.allclose(current_state.params, momentum05_final)

# Check that the two momentum formulas are equivalent with a constant learning rate
g_fn = jax.grad(lambda x: (x - 3.0) ** 2)
eta, mu = 0.1, 0.9
x1, v1 = 0.0, 0.0      # Sutskever form
x2, v2 = 0.0, 0.0      # PyTorch / Optax form
for _ in range(20):
    g1 = g_fn(x1); v1 = mu * v1 + eta * g1; x1 = x1 - v1
    g2 = g_fn(x2); v2 = mu * v2 + g2;       x2 = x2 - eta * v2
print(f"Sutskever form: {float(x1):.6f}   PyTorch/Optax form: {float(x2):.6f}")
assert jnp.allclose(x1, x2, atol=1e-5)

# ---------- Generic optimizers: interface modeled on Optax ----------
class Optimizer(NamedTuple):
    init: callable
    update: callable

def apply_updates(params, updates):
    return jax.tree.map(lambda p, u: p + u, params, updates)

class MomentumState(NamedTuple):
    velocity: dict

def sgd_momentum(lr, momentum=0.9):
    def init(params):
        return MomentumState(velocity=jax.tree.map(jnp.zeros_like, params))
    def update(grads, state, params=None):
        v = jax.tree.map(lambda v, g: momentum * v + g, state.velocity, grads)
        updates = jax.tree.map(lambda v: -lr * v, v)
        return updates, MomentumState(velocity=v)
    return Optimizer(init, update)

class AdamState(NamedTuple):
    count: jax.Array
    mu: dict
    nu: dict

def adam(lr, b1=0.9, b2=0.999, eps=1e-8):
    def init(params):
        # Note: mu and nu must be two independent sets of arrays; they can't share the same zeros object (see the donate_argnums Extra)
        return AdamState(count=jnp.zeros([], jnp.int32),
                         mu=jax.tree.map(jnp.zeros_like, params),
                         nu=jax.tree.map(jnp.zeros_like, params))
    def update(grads, state, params=None):
        t = state.count + 1
        mu = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g, state.mu, grads)
        nu = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g ** 2, state.nu, grads)
        bc1 = 1 - b1 ** t                       # bias correction
        bc2 = 1 - b2 ** t
        updates = jax.tree.map(lambda m, v: -lr * (m / bc1) / (jnp.sqrt(v / bc2) + eps), mu, nu)
        return updates, AdamState(count=t, mu=mu, nu=nu)
    return Optimizer(init, update)

# ---------- Model and data ----------
def init_mlp(key, sizes=(1, 32, 32, 1)):
    params = []
    keys = jax.random.split(key, len(sizes) - 1)
    for k, (n_in, n_out) in zip(keys, zip(sizes[:-1], sizes[1:])):
        params.append({'w': jax.random.normal(k, (n_in, n_out)) * jnp.sqrt(2.0 / n_in),
                       'b': jnp.zeros(n_out)})
    return params

def mlp(params, x):
    for layer in params[:-1]:
        x = jnp.tanh(x @ layer['w'] + layer['b'])
    return x @ params[-1]['w'] + params[-1]['b']

def loss_fn(params, x, y):
    return jnp.mean((mlp(params, x) - y) ** 2)

key = jax.random.key(0)
key, k_data, k_init = jax.random.split(key, 3)
X = jax.random.uniform(k_data, (256, 1), minval=-3.0, maxval=3.0)
Y = jnp.sin(X)
params0 = init_mlp(k_init)
print("Parameter structure:", jax.tree.map(jnp.shape, params0))

import numpy as np

# Self-check on a quadratic: the hand-written Adam / momentum SGD match a NumPy reference computed by hand, step by step
A = np.array([1.0, 4.0, 0.25], dtype=np.float32)        # different curvature in each dimension
C = np.array([3.0, -1.0, 2.0], dtype=np.float32)        # the optimum

def quad_loss(p):                                        # f(x) = 1/2 Σ a_i (x_i - c_i)^2
    return 0.5 * jnp.sum(A * (p['x'] - C) ** 2)

def numpy_adam_reference(x, steps, lr=0.1, b1=0.9, b2=0.999, eps=1e-8):
    m = np.zeros_like(x); v = np.zeros_like(x); traj = []
    for t in range(1, steps + 1):
        g = A * (x - C)                                  # gradient computed by hand
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
        m_hat = m / (1 - b1 ** t); v_hat = v / (1 - b2 ** t)
        x = x - lr * m_hat / (np.sqrt(v_hat) + eps)
        traj.append(x.copy())
    return traj

def numpy_momentum_reference(x, steps, lr=0.1, mu=0.9):
    v = np.zeros_like(x); traj = []
    for _ in range(steps):
        g = A * (x - C)
        v = mu * v + g
        x = x - lr * v
        traj.append(x.copy())
    return traj

def run_jax_optimizer(opt, steps):
    p = {'x': jnp.zeros(3)}
    state = opt.init(p); traj = []
    for _ in range(steps):
        g = jax.grad(quad_loss)(p)
        updates, state = opt.update(g, state, p)
        p = apply_updates(p, updates)
        traj.append(np.asarray(p['x']))
    return traj, state

x0 = np.zeros(3, dtype=np.float32)
adam_traj, adam_state = run_jax_optimizer(adam(0.1), 3)
for t, (ours, ref) in enumerate(zip(adam_traj, numpy_adam_reference(x0, 3)), start=1):
    print(f"Adam step {t}  JAX: {ours}  NumPy: {ref}")
    assert np.allclose(ours, ref, rtol=1e-5, atol=1e-6)   # differences of about 1e-6 come from float32 rounding in 1-β2^t
assert int(adam_state.count) == 3
# The effect of bias correction: at step 1 every coordinate moves by about lr=0.1, regardless of the gradient size
assert np.allclose(np.abs(adam_traj[0]), 0.1, atol=1e-6)

mom_traj, _ = run_jax_optimizer(sgd_momentum(0.1, momentum=0.9), 3)
for ours, ref in zip(mom_traj, numpy_momentum_reference(x0, 3)):
    assert np.allclose(ours, ref, rtol=1e-5, atol=1e-6)
print("The first 3 steps of Adam and of momentum SGD are allclose to the NumPy reference ✓")

def make_train(optimizer, num_steps):
    """Return a jit-compiled training function. Parameters and data are passed in as function arguments (no large arrays captured by closures);
    (params, opt_state) is the scan carry."""
    @jax.jit
    def train(params, x, y):
        def train_step(carry, _):
            params, opt_state = carry
            loss, grads = jax.value_and_grad(loss_fn)(params, x, y)   # x, y are arguments of train (tracers)
            updates, opt_state = optimizer.update(grads, opt_state, params)
            params = apply_updates(params, updates)
            return (params, opt_state), loss          # carry = state; y = per-step loss

        opt_state = optimizer.init(params)
        (params, opt_state), losses = lax.scan(train_step, (params, opt_state), None, length=num_steps)
        return params, opt_state, losses
    return train

results = {}
for name, opt in [("SGD", sgd_momentum(0.01, momentum=0.0)),
                  ("SGD+momentum", sgd_momentum(0.01, momentum=0.9)),
                  ("Adam", adam(0.01))]:
    p, s, losses = make_train(opt, 1000)(params0, X, Y)
    results[name] = losses
    print(f"{name:12s} initial loss {losses[0]:.4f} → final loss {losses[-1]:.5f}")

print("Step count in the Adam state:", s.count, "; first-moment structure matches the parameters:",
      jax.tree.structure(s.mu) == jax.tree.structure(p))

import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['PingFang SC', 'Heiti SC', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

fig, ax = plt.subplots(figsize=(6, 3.5))
for name, losses in results.items():
    ax.plot(losses, label=name)
ax.set_yscale('log'); ax.set_xlabel('step'); ax.set_ylabel('MSE loss (log scale)')
ax.set_title('Same model, three kinds of optimizer state'); ax.legend()
plt.tight_layout(); plt.show()

opt = adam(1e-2)

@partial(jax.jit, donate_argnums=(0, 1))     # donate params and opt_state; don't donate the data x, y
def train_step_donated(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = opt.update(grads, opt_state, params)
    return apply_updates(params, updates), opt_state, loss

params = jax.tree.map(jnp.copy, params0)
opt_state = opt.init(params)
old_leaf = params[0]['w']

params, opt_state, loss = train_step_donated(params, opt_state, X, Y)   # correct: overwrite the old references
print("Old parameter buffer has been deleted:", old_leaf.is_deleted())
assert old_leaf.is_deleted()
deleted_error = None
try:
    print(old_leaf + 1)
except RuntimeError as e:
    deleted_error = e
    print("Accessing the old array again → RuntimeError:", e)
assert deleted_error is not None
assert not X.is_deleted()                    # data that wasn't donated is still usable

for _ in range(200):
    params, opt_state, loss = train_step_donated(params, opt_state, X, Y)
print(f"Loss after 200 steps: {loss:.5f}")

# A donation that can't be used: no output array of the same shape → a warning, not an error
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    jax.jit(lambda a: a.sum(), donate_argnums=0)(jnp.ones(10))
msgs = [str(m.message) for m in w]
print("Warnings:", [m[:60] for m in msgs])
assert any("donated buffers were not usable" in m for m in msgs)

# ❌ Wrong (same buffer donated twice): Adam's mu and nu share the same zeros object
def adam_init_shared(params):
    zeros = jax.tree.map(jnp.zeros_like, params)
    return AdamState(count=jnp.zeros([], jnp.int32), mu=zeros, nu=zeros)   # mu and nu are the very same arrays

params = jax.tree.map(jnp.copy, params0)
bad_opt_state = adam_init_shared(params)
double_donate_error = None
try:
    train_step_donated(params, bad_opt_state, X, Y)
except Exception as e:
    double_donate_error = e
    print(f"Actual result → {type(e).__name__}: {str(e)[:120]}")
assert double_donate_error is not None and "donate the same buffer twice" in str(double_donate_error)
# (XLA also prints a C++ log line with the same content to stderr; you can ignore it)

# ✅ Fix: create mu and nu separately (the way adam(...).init above does it)
params = jax.tree.map(jnp.copy, params0)
good_opt_state = opt.init(params)
assert good_opt_state.mu[0]['w'] is not good_opt_state.nu[0]['w']
params, good_opt_state, loss = train_step_donated(params, good_opt_state, X, Y)
print(f"Independent buffers + donation: runs fine, loss = {loss:.5f}")

def stateful_counter(count_state, increment):
    """Take the counter state and an increment, and return the new state."""
    new_count = count_state['count'] + increment
    return {'count': new_count}

initial_state = {'count': jnp.array(0)}

# Without jit
state = stateful_counter(initial_state, 1)
state = stateful_counter(state, 2)
print("Result of manual calls:", state)

# With jit
jit_counter = jax.jit(stateful_counter)
jit_state = jit_counter(initial_state, 1)
jit_state = jit_counter(jit_state, 2)
print("Result of JIT calls:", jit_state)

# ❌ Old style (outdated): the top-level alias jax.tree_map was removed in JAX 0.12
tree_map_error = None
try:
    assert jax.tree_util.tree_all(jax.tree_map(lambda x, y: jnp.array_equal(x, y), state, jit_state))
except AttributeError as e:
    tree_map_error = e
    print("Actual result → AttributeError:", e)
assert tree_map_error is not None

# ✅ Recommended
assert jax.tree.all(jax.tree.map(lambda x, y: jnp.array_equal(x, y), state, jit_state))
print("jax.tree.map assertion passed: both results match")

traces = []
@jax.jit
def bump(state):
    traces.append(jax.tree.structure(state))     # recorded during tracing
    return jax.tree.map(lambda x: x + 1, state)

bump({'count': jnp.array(0)})
bump({'count': jnp.array(5)})                     # same structure: cache hit
bump({'count': jnp.array(0), 'steps': jnp.array(0)})   # new key: retrace
print("Number of traces:", len(traces))
for t in traces: print("  ", t)
assert len(traces) == 2

def predict_and_update(params, x):
    """Forward prediction; the parameter state is returned as is (not modified here)."""
    pred = params['w'] * x + params['b']
    return pred, params

def loss_fn_linear(params, x, y_target):
    pred, _ = predict_and_update(params, x)
    return jnp.mean((pred - y_target) ** 2)

# Data: the ideal parameters are w=2.5, b=1.0
x_data = jnp.array([1.0, 2.0, 3.0])
y_data = jnp.array([3.5, 5.5, 7.5])
params = {'w': jnp.array(2.0), 'b': jnp.array(1.0)}

grad_fn_linear = jax.grad(loss_fn_linear, argnums=0)
grads = grad_fn_linear(params, x_data, y_data)
print("Gradient (same structure as params):", grads)

loss_val, grads = jax.value_and_grad(loss_fn_linear)(params, x_data, y_data)
print("Loss:", loss_val, " gradient:", grads)