Chapter 10: Random Numbers and Sharding

# ⚠️ Must be set before import jax: simulate 8 devices on CPU (if jax is already imported, restart the kernel)
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import sys, re, io, json, time, socket, subprocess, warnings
import random as py_random            # Python's standard random module, aliased so it isn't confused with jax.random
from collections import Counter

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.sharding import NamedSharding, PartitionSpec as P
from scipy import stats
import flax
from flax import nnx

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| Flax", flax.__version__)
print("Devices", jax.devices())
assert jax.__version__ == "0.12.0.1"
assert jax.device_count() == 8


def bench(f, *args, warmup=2, runs=10):
    """Timing: warm up `warmup` times (including compilation), then run `runs` times and return the median in ms. For printing only."""
    for _ in range(warmup):
        jax.block_until_ready(f(*args))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args))
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts)) * 1e3


COLLECTIVE_OPS = ["all-reduce", "all-gather", "reduce-scatter", "all-to-all", "collective-permute"]

def collectives(hlo_text):
    """Count how many times each kind of collective appears in compiled HLO text (including the async *-start forms)."""
    return {op: len(re.findall(op + r"(?:-start)?\(", hlo_text)) for op in COLLECTIVE_OPS}


def run_py(code, *args, env=None, timeout=600):
    """Run code in a fresh Python subprocess and return the object parsed from its last 'JSON:' output line."""
    full_env = dict(os.environ, **(env or {}))
    r = subprocess.run([sys.executable, "-c", code, *map(str, args)],
                       capture_output=True, text=True, env=full_env, timeout=timeout)
    assert r.returncode == 0, r.stderr[-3000:]
    lines = [l for l in r.stdout.splitlines() if l.startswith("JSON:")]
    assert lines, r.stdout[-2000:] + r.stderr[-2000:]
    return json.loads(lines[-1][5:])


def err_line(e):
    """Exception type + first line of the message (used when demonstrating wrong code)."""
    msg = str(e).strip().splitlines()[0] if str(e).strip() else ""
    return f"{type(e).__name__}: {msg[:170]}"

print("Common helpers defined")

# NumPy's global state: same seed, we only swap the call order of two functions
def init_weights_np():
    return np.random.normal(size=3)

def sample_noise_np():
    return np.random.normal(size=3)

np.random.seed(0); w_a = init_weights_np(); n_a = sample_noise_np()
np.random.seed(0); n_b = sample_noise_np(); w_b = init_weights_np()   # reversed order
print("Weights with init first, then noise:", w_a)
print("Weights with noise first, then init:", w_b)

assert not np.array_equal(w_a, w_b)       # the weights changed
assert np.array_equal(w_a, n_b)            # and "weights" and "noise" swapped identities
print("✓ NumPy global state: results depend on call order")

# JAX: random numbers depend only on the key, not on call order
def init_weights(key):
    return jax.random.normal(key, (3,))

def sample_noise(key):
    return jax.random.normal(key, (3,))

k_w, k_n = jax.random.split(jax.random.key(0))
w1 = init_weights(k_w); n1 = sample_noise(k_n)
n2 = sample_noise(k_n); w2 = init_weights(k_w)      # reversed order
assert jnp.array_equal(w1, w2) and jnp.array_equal(n1, n2)

# The key is a regular argument: jit traces / compiles once, and changing the key doesn't recompile
trace_count = 0
def noisy(key, x):
    global trace_count
    trace_count += 1                  # runs only during tracing
    return x + jax.random.normal(key, x.shape)

noisy_jit = jax.jit(noisy)
x0 = jnp.zeros(3)
outs = [noisy_jit(k, x0) for k in jax.random.split(jax.random.key(1), 4)]
print("Results for 4 different keys are all distinct:", len({tuple(np.asarray(o).tolist()) for o in outs}) == 4,
      "| trace count =", trace_count)
assert trace_count == 1
assert len({tuple(np.asarray(o).tolist()) for o in outs}) == 4
print("✓ JAX: order-independent and reproducible; a key as an argument doesn't trigger recompilation")

k = jax.random.key(0)
raw = jax.random.PRNGKey(0)

# The standard way to check "is this a typed key": jax.dtypes.prng_key is the abstract parent of all key dtypes
print("key(0).dtype =", k.dtype, "| is prng_key:", jnp.issubdtype(k.dtype, jax.dtypes.prng_key))
print("PRNGKey(0).dtype =", raw.dtype, "| is prng_key:", jnp.issubdtype(raw.dtype, jax.dtypes.prng_key))
print("The key remembers its algorithm:", jax.random.key_impl(k))
assert jnp.issubdtype(k.dtype, jax.dtypes.prng_key)
assert not jnp.issubdtype(raw.dtype, jax.dtypes.prng_key)
assert str(jax.random.key_impl(k)) == "threefry2x32"

# Shape semantics: key.shape is the number of keys; key_data adds one trailing dimension
print(f"\n{'Expression':28s} {'key.shape':>10s} {'key_data.shape':>15s}")
for name, kk in [("key(0)", k),
                 ("split(key, 4)", jax.random.split(k, 4)),
                 ("split(key, (2, 3))", jax.random.split(k, (2, 3))),
                 ("key(0, impl='rbg')", jax.random.key(0, impl="rbg")),
                 ("split(rbg_key, (2, 3))", jax.random.split(jax.random.key(0, impl="rbg"), (2, 3))),
                 ("key(0, impl='philox2x32')", jax.random.key(0, impl="philox2x32"))]:
    print(f"{name:28s} {str(kk.shape):>10s} {str(jax.random.key_data(kk).shape):>15s}")

assert jax.random.split(k, (2, 3)).shape == (2, 3)
assert jax.random.key_data(jax.random.split(k, (2, 3))).shape == (2, 3, 2)
assert jax.random.key_data(jax.random.split(jax.random.key(0, impl="rbg"), (2, 3))).shape == (2, 3, 4)
assert jax.random.key_data(jax.random.key(0, impl="philox2x32")).shape == (1,)

ks = jax.random.split(k, 4)
ok_ops = {
    "ks[1:3] (slicing)":          lambda: ks[1:3].shape,
    "ks.reshape(2, 2)":           lambda: ks.reshape(2, 2).shape,
    "ks.reshape(2, 2).T":         lambda: ks.reshape(2, 2).T.shape,
    "jnp.stack([k, k])":          lambda: jnp.stack([k, k]).shape,
    "jnp.concatenate([ks, ks])":  lambda: jnp.concatenate([ks, ks]).shape,
    "jnp.where(mask, ks, ks[::-1])": lambda: jnp.where(jnp.array([True, False, True, False]), ks, ks[::-1]).shape,
    "jnp.broadcast_to(k, (3,))":  lambda: jnp.broadcast_to(k, (3,)).shape,
    "ks == ks (comparison)":      lambda: (ks == ks).tolist(),
}
for name, f in ok_ops.items():
    print(f"✅ {name:32s} → {f()}")
assert jnp.stack([k, k]).dtype == k.dtype and bool(jnp.all(ks == ks))

# ❌ Wrong: arithmetic / reductions / type conversions on typed keys (trying to "derive" a new key by hand, or save a key as data)
bad_ops = {
    "k + 1":                     lambda: k + 1,
    "k * 2":                     lambda: k * 2,
    "k ^ k":                     lambda: k ^ k,
    "jnp.sum(ks)":               lambda: jnp.sum(ks),
    "jnp.unique(ks)":            lambda: jnp.unique(ks),
    "jnp.asarray(k, jnp.uint32)": lambda: jnp.asarray(k, jnp.uint32),
    "np.asarray(k)":             lambda: np.asarray(k),
}
bad_errors = {}
for name, f in bad_ops.items():
    try:
        out = f()
        print(f"(no error) {name} → {out}")
    except Exception as e:
        bad_errors[name] = type(e).__name__
        print(f"❌ {name:28s} → {err_line(e)}")

assert set(bad_errors) == set(bad_ops)          # all of them raise
assert bad_errors["k + 1"] == "TypeError" and bad_errors["np.asarray(k)"] == "TypeError"
# For comparison: doing the same with a raw uint32 key "silently succeeds", which is exactly the danger of the old style
print("\nRaw key: PRNGKey(0) + 1 =", raw + 1, "(silently succeeds)")
assert jnp.array_equal(raw + 1, jnp.array([1, 1], jnp.uint32))

# ✅ Correct: derive new keys with fold_in / split; get the underlying bits with key_data
k_step1 = jax.random.fold_in(k, 1)                  # instead of "k + 1"
raw_bits = jax.random.key_data(k)                   # instead of "np.asarray(k)"
print("Underlying data of fold_in(k, 1):", jax.random.key_data(k_step1))
print("key_data(k) =", raw_bits, raw_bits.dtype)
assert k_step1.dtype == k.dtype and not bool(k_step1 == k)
assert raw_bits.dtype == jnp.uint32 and raw_bits.shape == (2,)

# ❌ Wrong: save only the raw data of an rbg key and restore it without impl (interpreted as the default threefry2x32)
rbg_key = jax.random.key(7, impl="rbg")
rbg_raw = jax.random.key_data(rbg_key)
try:
    restored = jax.random.wrap_key_data(rbg_raw)
    print("(no error) restored as", restored.dtype)
    wrap_err = None
except Exception as e:
    wrap_err = e
    print("❌", err_line(e))
assert isinstance(wrap_err, TypeError)

# ✅ Correct: save (impl name, key_data), restore with wrap_key_data(..., impl=impl)
def save_keys(keys):
    """Serialize a typed key array to bytes: impl name + raw uint32 data."""
    buf = io.BytesIO()
    np.savez(buf, impl=str(jax.random.key_impl(keys)), data=np.asarray(jax.random.key_data(keys)))
    return buf.getvalue()

def load_keys(blob):
    z = np.load(io.BytesIO(blob))
    return jax.random.wrap_key_data(jnp.asarray(z["data"]), impl=str(z["impl"]))

for impl in ["threefry2x32", "rbg", "philox2x32"]:
    keys = jax.random.split(jax.random.key(7, impl=impl), (2, 3))
    back = load_keys(save_keys(keys))
    same = bool(jnp.all(back == keys)) and back.dtype == keys.dtype
    same_samples = jnp.array_equal(jax.random.normal(back[1, 2], (4,)), jax.random.normal(keys[1, 2], (4,)))
    print(f"{impl:13s} after round trip dtype={back.dtype}, shape={back.shape}, same key={same}, same samples={bool(same_samples)}")
    assert same and same_samples

# (1) Seed truncation
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    same_key = bool(jax.random.key(2**32 + 5) == jax.random.key(5))
print("32-bit mode: key(2**32 + 5) == key(5)?", same_key, "| warnings:", len(caught))
assert same_key and len(caught) == 0

x64_res = run_py("""
import json, jax
jax.config.update("jax_enable_x64", True)
a = jax.random.key(2**32 + 5)
print("JSON:" + json.dumps({"data": jax.random.key_data(a).tolist(), "same": bool(a == jax.random.key(5))}))
""")
print("x64 mode (subprocess): key_data(key(2**32 + 5)) =", x64_res["data"], "| same as key(5)?", x64_res["same"])
assert x64_res["data"] == [1, 5] and not x64_res["same"]

# A safe way to use a large integer seed: split it into high and low halves and fold_in each; the result doesn't depend on the x64 setting
def key_from_big_seed(seed: int):
    hi, lo = (seed >> 32) & 0xFFFFFFFF, seed & 0xFFFFFFFF
    return jax.random.fold_in(jax.random.fold_in(jax.random.key(0), hi), lo)
assert not bool(key_from_big_seed(2**32 + 5) == key_from_big_seed(5))

# (2) jax_legacy_prng_key
print("\nDefault of jax_legacy_prng_key:", jax.config.jax_legacy_prng_key)
with jax.legacy_prng_key("error"):
    try:
        jax.random.normal(jax.random.PRNGKey(0), (2,))
        legacy_err = None
    except Exception as e:
        legacy_err = e
        print("Passing a raw key in 'error' mode →", err_line(e))
    ok_typed = jax.random.normal(jax.random.key(0), (2,))      # typed keys aren't affected
assert isinstance(legacy_err, ValueError) and ok_typed.shape == (2,)

IMPLS = ["threefry2x32", "threefry4x32", "philox2x32", "philox4x32", "rbg", "unsafe_rbg"]
samples = {}
print(f"{'impl':13s} {'dtype':10s} {'key_data':>9s}  normal(key(0), (3,))")
for impl in IMPLS:
    kk = jax.random.key(0, impl=impl)
    samples[impl] = np.asarray(jax.random.normal(kk, (3,)))
    print(f"{impl:13s} {str(kk.dtype):10s} {str(jax.random.key_data(kk).shape):>9s}  {samples[impl]}")

expected_dtype = {"threefry2x32": "key<fry>", "threefry4x32": "key<fry4>", "philox2x32": "key<phx2>",
                  "philox4x32": "key<phx4>", "rbg": "key<rbg>", "unsafe_rbg": "key<urbg>"}
assert all(str(jax.random.key(0, impl=i).dtype) == d for i, d in expected_dtype.items())
# Same seed, different algorithms give different random numbers (rbg and unsafe_rbg share the "generation" part; only key derivation differs)
distinct = {tuple(v.tolist()) for v in samples.values()}
assert len(distinct) == 5 and np.array_equal(samples["rbg"], samples["unsafe_rbg"])

# Global default algorithm: environment variable (subprocess, effective from startup) and context manager (this process, temporary)
env_res = run_py("""
import json, jax
print("JSON:" + json.dumps({
    "config": jax.config.jax_default_prng_impl,
    "key_dtype": str(jax.random.key(0).dtype),
    "raw_shape": list(jax.random.PRNGKey(0).shape),                      # raw keys are interpreted using the global config
    "explicit": str(jax.random.key(0, impl="threefry2x32").dtype)}))     # an explicit impl isn't affected
""", env={"JAX_DEFAULT_PRNG_IMPL": "rbg"})
print("Subprocess with JAX_DEFAULT_PRNG_IMPL=rbg:", env_res)
assert env_res == {"config": "rbg", "key_dtype": "key<rbg>", "raw_shape": [4], "explicit": "key<fry>"}

with jax.default_prng_impl("philox4x32"):
    inside = str(jax.random.key(0).dtype)
outside = str(jax.random.key(0).dtype)
print("Inside with jax.default_prng_impl('philox4x32'):", inside, "| after exiting:", outside)
assert inside == "key<phx4>" and outside == "key<fry>"

print(f"{'impl':13s} {'gen 4M normal':>14s} {'split 100k':>11s} {'1st compile':>9s}")
for impl in IMPLS:
    kk = jax.random.key(0, impl=impl)
    gen = jax.jit(lambda key: jax.random.normal(key, (4_000_000,)))
    spl = jax.jit(lambda key: jax.random.split(key, 100_000))
    t0 = time.perf_counter(); gen.lower(kk).compile(); t_compile = (time.perf_counter() - t0) * 1e3
    print(f"{impl:13s} {bench(gen, kk):11.2f} ms {bench(spl, kk):8.2f} ms {t_compile:6.0f} ms")

kr = jax.random.key(0, impl="rbg")
ku = jax.random.key(0, impl="unsafe_rbg")
halves = jax.random.key_data(kr).reshape(2, 2)          # rbg key = two threefry keys concatenated
tf_split = jnp.concatenate([jax.random.key_data(jax.random.split(jax.random.wrap_key_data(h), 3)) for h in halves], -1)

rbg_like_tf = bool(jnp.array_equal(jax.random.key_data(jax.random.split(kr, 3)), tf_split))
urbg_like_tf = bool(jnp.array_equal(jax.random.key_data(jax.random.split(ku, 3)), tf_split))
print("rbg's split == threefry split on each half:", rbg_like_tf)
print("unsafe_rbg's split == same as above:", urbg_like_tf, "(it uses RBG as a hash; quality untested)")
assert rbg_like_tf and not urbg_like_tf

# ❌ Wrong: with rbg, assuming "vmap equals calling one at a time" (e.g. cross-checking a vmap version against a loop version)
keys_rbg = jax.random.split(jax.random.key(0, impl="rbg"), 8)
try:
    v = jax.vmap(lambda kk: jax.random.normal(kk))(keys_rbg)
    loop = jnp.stack([jax.random.normal(kk) for kk in keys_rbg])
    first_only = jax.random.normal(keys_rbg[0], (8,))
    print("No error, but vmap == one at a time?", bool(jnp.array_equal(v, loop)),
          "| vmap == normal(keys[0], (8,))?", bool(jnp.array_equal(v, first_only)))
except Exception as e:
    print("❌", err_line(e))
assert not jnp.array_equal(v, loop) and jnp.array_equal(v, first_only)

# ✅ Correct: when you need "vmap matches calling one at a time", use threefry / philox (the default threefry2x32 is fine)
for impl in ["threefry2x32", "philox4x32"]:
    keys_ok = jax.random.split(jax.random.key(0, impl=impl), 8)
    v_ok = jax.vmap(lambda kk: jax.random.normal(kk))(keys_ok)
    loop_ok = jnp.stack([jax.random.normal(kk) for kk in keys_ok])
    print(f"{impl:13s} vmap == one at a time: {bool(jnp.array_equal(v_ok, loop_ok))}")
    assert jnp.array_equal(v_ok, loop_ok)

print("jax_threefry_partitionable =", jax.config.jax_threefry_partitionable)
assert jax.config.jax_threefry_partitionable is True
# The jax.random module docs state which version made it the default
assert "As of JAX v.0.5.0, this is the" in jax.random.__doc__
print("jax.random docs: 'As of JAX v.0.5.0, this is the default.' ✓")

# Three ways to change it: jax.config.update / environment variable JAX_THREEFRY_PARTITIONABLE / context manager
env_flag = run_py('import json, jax; print("JSON:" + json.dumps(jax.config.jax_threefry_partitionable))',
                  env={"JAX_THREEFRY_PARTITIONABLE": "0"})
with jax.threefry_partitionable(False):
    ctx_flag = jax.config.jax_threefry_partitionable
print("Env var JAX_THREEFRY_PARTITIONABLE=0 →", env_flag, "| inside context manager →", ctx_flag,
      "| after exiting →", jax.config.jax_threefry_partitionable)
assert env_flag is False and ctx_flag is False and jax.config.jax_threefry_partitionable is True

FLAG_SCRIPT = """
import sys, json, jax, jax.numpy as jnp
jax.config.update("jax_threefry_partitionable", sys.argv[1] == "1")
k = jax.random.key(0)
x10 = jax.random.normal(k, (10,))
print("JSON:" + json.dumps({
    "flag": jax.config.jax_threefry_partitionable,
    "normal_8": jax.random.normal(k, (8,)).tolist(),
    "uniform_3": jax.random.uniform(k, (3,)).tolist(),
    "split": jax.random.key_data(jax.random.split(k)).tolist(),
    "fold_in": jax.random.key_data(jax.random.fold_in(k, 1)).tolist(),
    "dot10": float(jnp.dot(x10, x10)),      # the number from the install check script in Chapter 1, Section 1.4
}))
"""
flag_runs = {flag: [run_py(FLAG_SCRIPT, int(flag)) for _ in range(2)] for flag in (True, False)}
on, off = flag_runs[True][0], flag_runs[False][0]
for name in ["normal_8", "uniform_3", "split", "fold_in", "dot10"]:
    print(f"{name:10s} on:  {on[name]}\n{'':10s} off: {off[name]}\n")

assert on["flag"] is True and off["flag"] is False
assert flag_runs[True][0] == flag_runs[True][1] and flag_runs[False][0] == flag_runs[False][1]   # reproducible under the same setting
assert on["normal_8"] != off["normal_8"] and on["uniform_3"] != off["uniform_3"]               # random bits differ
assert on["split"] != off["split"]                                                              # split differs too
assert on["fold_in"] == off["fold_in"]                                                          # fold_in isn't affected
# Matches the measured values in Chapter 1, Section 1.4: on 9.7387…, off 2.6585…
assert np.isclose(on["dot10"], 9.738718032836914) and np.isclose(off["dot10"], 2.6585559844970703)
# This process's default setting is bit-identical to the "on" subprocess
assert jax.random.normal(jax.random.key(0), (8,)).tolist() == on["normal_8"]
print("✓ On vs. off: random bits and split differ, fold_in is the same; each setting is reproducible")

from jax.extend.random import threefry_2x32, threefry2x32_p

kd = jax.random.key_data(k)
n = 8
with jax.threefry_partitionable(False):
    bits_off = jax.random.bits(k, (n,))
    manual_off = threefry_2x32(kd, jnp.arange(n, dtype=jnp.uint32))   # internally splits the counter in two halves
with jax.threefry_partitionable(True):
    bits_on = jax.random.bits(k, (n,))
    w0, w1 = threefry2x32_p.bind(kd[0], kd[1], jnp.zeros(n, jnp.uint32), jnp.arange(n, dtype=jnp.uint32))
    manual_on = w0 ^ w1                                                  # counter = global index i, XOR the two words
print("off: bits        =", bits_off)
print("off: by hand     =", manual_off)
print("on:  bits        =", bits_on)
print("on:  by hand     =", manual_on)
assert jnp.array_equal(bits_off, manual_off) and jnp.array_equal(bits_on, manual_on)

for flag in (True, False):
    with jax.threefry_partitionable(flag):
        prefix_1d = bool(jnp.array_equal(jax.random.normal(k, (10,))[:5], jax.random.normal(k, (5,))))
        prefix_2d = bool(jnp.array_equal(jax.random.normal(k, (4, 5))[:2], jax.random.normal(k, (2, 5))))
        split_prefix = bool(jnp.all(jax.random.split(k, 5)[:3] == jax.random.split(k, 3)))
        split_is_fold = bool(jnp.all(jax.random.split(k, 4) == jnp.stack([jax.random.fold_in(k, i) for i in range(4)])))
    print(f"switch={flag!s:5}: prefix property 1D={prefix_1d}, 2D row prefix={prefix_2d}, "
          f"split(k,5)[:3]==split(k,3): {split_prefix}, split(k,n)[i]==fold_in(k,i): {split_is_fold}")
    assert prefix_1d == prefix_2d == split_prefix == split_is_fold == flag

mesh = jax.make_mesh((8,), ("x",))
N = 1024
k_sh = jax.random.key(0)

def threefry_values_per_device(hlo_text, flag):
    """Estimate how many random values each device generates from the xor shape in the threefry round function in the HLO."""
    shapes = re.findall(r"u32\[([0-9,]+)\][^\n]*? xor\(", hlo_text)
    elems = max(int(np.prod([int(d) for d in s.split(",")])) for s in shapes)
    return elems * (1 if flag else 2)

rows = []
for flag in (True, False):
    with jax.threefry_partitionable(flag):
        ref = jax.jit(lambda key: jax.random.normal(key, (8, N)))(k_sh)          # unsharded reference
        for spec in (P("x"), P(None, "x")):
            f = jax.jit(lambda key: jax.random.normal(key, (8, N)), out_shardings=NamedSharding(mesh, spec))
            out = f(k_sh)
            txt = f.lower(k_sh).compile().as_text()
            rows.append(dict(flag=flag, spec=str(spec), equal=bool(jnp.array_equal(out, ref)),
                             comm=collectives(txt), per_dev=threefry_values_per_device(txt, flag),
                             shard_shape=out.addressable_shards[0].data.shape))

print(f"{'flag':5s} {'sharding':14s} {'= unsharded':>10s} {'shard/dev':>10s} {'gen/dev':>9s}  collectives")
for r in rows:
    comm = {kk: v for kk, v in r["comm"].items() if v} or "none"
    print(f"{r['flag']!s:5s} {r['spec']:14s} {r['equal']!s:>10s} {str(r['shard_shape']):>12s} {r['per_dev']:>9d}  {comm}")

R = {(r["flag"], r["spec"]): r for r in rows}
assert all(r["equal"] for r in rows)                                     # all 4 cases match the unsharded result
# On: no collectives with either split, and each device generates only its own 8*N/8 values
assert all(sum(R[(True, s)]["comm"].values()) == 0 and R[(True, s)]["per_dev"] == N for s in ("P('x',)", "P(None, 'x')"))
# Off + split by rows: collectives (all-gather / collective-permute) are needed to move paired elements to the right device
assert sum(R[(False, "P('x',)")]["comm"].values()) > 0
# Off + split by columns: no communication, but every device computes the whole 8*N array (8x redundant computation)
assert sum(R[(False, "P(None, 'x')")]["comm"].values()) == 0 and R[(False, "P(None, 'x')")]["per_dev"] == 8 * N
print("\n✓ On: generated locally, zero communication; off: either communication, or every device redundantly computes the whole array")

SEARCH_SCRIPT = """
import sys, json, jax, jax.numpy as jnp, numpy as np
from jax import lax
jax.config.update("jax_threefry_partitionable", sys.argv[1] == "1")
jax.config.update("jax_enable_x64", sys.argv[2] == "1")
TARGET = 10.811179161071777
out = {}
for impl in ["threefry2x32", "threefry4x32", "philox2x32", "philox4x32", "rbg", "unsafe_rbg"]:
    def one(seed):
        x = jax.random.normal(jax.random.key(seed, impl=impl), (10,))
        return jnp.dot(x, x)
    vals = np.asarray(jax.jit(lambda s: lax.map(one, s))(jnp.arange(1000)), np.float64)
    d = np.abs(vals - TARGET)
    out[impl] = {"seed0": float(vals[0]), "min_dist": float(d.min()), "argmin": int(d.argmin())}
print("JSON:" + json.dumps(out))
"""
search = {(f, x): run_py(SEARCH_SCRIPT, f, x) for f in (1, 0) for x in (0, 1)}
best = min(((v["min_dist"], cfg, impl, v["argmin"]) for cfg, res in search.items() for impl, v in res.items()))
print(f"{'partitionable, x64':20s} {'dot(x,x) for threefry2x32 seed 0':>32s}")
for (f, x), res in search.items():
    print(f"{f!s:>10s}, {x!s:<8s} {res['threefry2x32']['seed0']:>32.6f}")
print(f"\nClosest to 10.811179 among all {len(search) * 6 * 1000} combinations: distance {best[0]:.2e}"
      f" (partitionable={best[1][0]}, x64={best[1][1]}, impl={best[2]}, seed={best[3]})")

assert np.isclose(search[(1, 0)]["threefry2x32"]["seed0"], 9.7387, atol=1e-3)     # 0.12 default
assert np.isclose(search[(0, 0)]["threefry2x32"]["seed0"], 2.6586, atol=1e-3)     # default before 0.5.0
assert best[0] > 1e-4          # no combination reproduces 10.811179 within floating-point error (float32 error is about 1e-6)

# ❌ Wrong: key isn't updated after split, so the next split starts from the same old key → you get the same child key twice
key = jax.random.key(0)
try:
    _, k_init = jax.random.split(key)          # the new key was thrown into _
    w_init = jax.random.normal(k_init, (4,))
    _, k_noise = jax.random.split(key)         # split from the old key again
    noise = jax.random.normal(k_noise, (4,))
    print("No error, but the \"init\" and the \"noise\" are identical:", bool(jnp.array_equal(w_init, noise)))
except Exception as e:
    print("❌", err_line(e))
assert jnp.array_equal(w_init, noise)

# ✅ Correct: after every split, assign the new key back to key and never use the old one again
key = jax.random.key(0)
key, k_init = jax.random.split(key)
w_init = jax.random.normal(k_init, (4,))
key, k_noise = jax.random.split(key)
noise = jax.random.normal(k_noise, (4,))
print("Init and noise differ:", not bool(jnp.array_equal(w_init, noise)))
assert not jnp.array_equal(w_init, noise) and not bool(k_init == k_noise)

# ❌ Wrong: derive the key for worker i with key(seed + i) → experiments with different seeds share random numbers
def worker_keys_bad(seed, n):
    return [jax.random.key(seed + i) for i in range(n)]
try:
    run_seed0 = worker_keys_bad(0, 4)
    run_seed1 = worker_keys_bad(1, 4)
    collide = bool(run_seed0[1] == run_seed1[0])
    print("No error, but seed0/worker1 and seed1/worker0 have the same key:", collide)
except Exception as e:
    print("❌", err_line(e))
assert collide
assert jnp.array_equal(jax.random.normal(run_seed0[1], (3,)), jax.random.normal(run_seed1[0], (3,)))

# ✅ Correct: split n keys at once from the experiment's main key
def worker_keys(seed, n):
    return jax.random.split(jax.random.key(seed), n)

a, b = worker_keys(0, 4), worker_keys(1, 4)
all_keys = [tuple(r) for r in np.asarray(jax.random.key_data(jnp.concatenate([a, b]))).tolist()]
print("8 keys across the two experiments, all distinct:", len(set(all_keys)) == 8)
assert len(set(all_keys)) == 8

# One split(key, n) vs. n splits in a Python loop (both results are usable, but the loop needs n dispatches)
def loop_split(key, n):
    out = []
    for _ in range(n):
        key, sub = jax.random.split(key)
        out.append(sub)
    return out
t_loop = bench(lambda kk: loop_split(kk, 200), jax.random.key(0), runs=5)      # bench warms up first, excluding the first compilation
t_once = bench(lambda kk: jax.random.split(kk, 200), jax.random.key(0))
print(f"200 keys: Python loop of splits {t_loop:.2f} ms, one split(key, 200) {t_once:.3f} ms (median after warmup, printed only)")

# ❌ Wrong: every step of the training loop generates the dropout mask from the same key → every step has the same mask
base = jax.random.key(42)
x_act = jnp.ones((4, 6))

@jax.jit
def dropout_mask(key, x):
    return jax.random.bernoulli(key, 0.5, x.shape)

try:
    masks_bad = [dropout_mask(base, x_act) for step in range(5)]
    print("No error, but the masks of all 5 steps are identical:", all(bool(jnp.array_equal(m, masks_bad[0])) for m in masks_bad))
except Exception as e:
    print("❌", err_line(e))
assert all(jnp.array_equal(m, masks_bad[0]) for m in masks_bad)

# ✅ Correct: inside jit, derive each step's key with fold_in(base, step), passing step as a dynamic argument
trace_count = 0
@jax.jit
def dropout_mask_at(base, step, x):
    global trace_count
    trace_count += 1
    return jax.random.bernoulli(jax.random.fold_in(base, step), 0.5, x.shape)

masks = [dropout_mask_at(base, jnp.uint32(s), x_act) for s in range(5)]
distinct_masks = len({np.asarray(m).tobytes() for m in masks})
print("Masks of the 5 steps are all distinct:", distinct_masks == 5, "| trace count:", trace_count)
assert distinct_masks == 5 and trace_count == 1
# Random access: recompute step 3 directly, without replaying earlier steps
assert jnp.array_equal(dropout_mask_at(base, jnp.uint32(3), x_act), masks[3])

# Derive from data indices: the same example gets the same augmentation noise no matter which batch it lands in
k_aug = jax.random.key(7)
def augment(idx, x):                          # x is a single example, idx is its index in the dataset
    return x + 0.1 * jax.random.normal(jax.random.fold_in(k_aug, idx), x.shape)
data = jnp.arange(10.0).reshape(10, 1) * jnp.ones((10, 3))
batch_a = jnp.array([5, 1, 7]); batch_b = jnp.array([2, 5, 9, 0])
aug_a = jax.vmap(augment)(batch_a, data[batch_a])
aug_b = jax.vmap(augment)(batch_b, data[batch_b])
print("Example 5 gets the same augmentation in two different batches:", bool(jnp.array_equal(aug_a[0], aug_b[1])))
assert jnp.array_equal(aug_a[0], aug_b[1])

try:
    jax.random.fold_in(base, -1)
except OverflowError as e:
    print("fold_in(base, -1) →", err_line(e))

# ❌ Wrong: the key is in the carry, but the loop body never updates it → every step gets the same random number
def body_bad(key, _):
    return key, jax.random.normal(key)             # passes the old key to the next step unchanged

try:
    _, ys_bad = lax.scan(body_bad, jax.random.key(0), None, length=6)
    print("No error, but the random numbers of the 6 steps:", ys_bad)
except Exception as e:
    print("❌", err_line(e))
assert bool(jnp.all(ys_bad == ys_bad[0]))

# ✅ Correct, option A: split every step and hand the new key to the next step as the carry
def body(key, _):
    key, sub = jax.random.split(key)
    return key, jax.random.normal(sub)
_, ys_a = jax.jit(lambda kk: lax.scan(body, kk, None, length=6))(jax.random.key(0))

# Bit-identical to the equivalent Python loop
kk, ys_loop = jax.random.key(0), []
for _ in range(6):
    kk, sub = jax.random.split(kk)
    ys_loop.append(jax.random.normal(sub))
assert jnp.array_equal(ys_a, jnp.stack(ys_loop))

# ✅ Correct, option B: split T keys ahead of time and feed them in as xs (no key in the carry)
_, ys_b = lax.scan(lambda c, kt: (c, jax.random.normal(kt)), None, jax.random.split(jax.random.key(0), 6))
print("Option A:", ys_a)
print("Option B:", ys_b)
assert len(set(np.asarray(ys_a).tolist())) == 6 and len(set(np.asarray(ys_b).tolist())) == 6

# ❌ Wrong: the key isn't mapped under vmap (in_axes=None) → every example gets the same "random numbers"
def noisy_sample(key, x):
    return x + jax.random.normal(key, x.shape)

xs_batch = jnp.zeros((4, 3))
try:
    out_bad = jax.vmap(noisy_sample, in_axes=(None, 0))(jax.random.key(0), xs_batch)
    print("No error, but the noise is identical for all 4 examples:", bool(jnp.all(out_bad == out_bad[0])))
except Exception as e:
    print("❌", err_line(e))
assert bool(jnp.all(out_bad == out_bad[0]))

# ✅ Correct: split B keys and map them along axis 0 together with the examples
keys_b = jax.random.split(jax.random.key(0), xs_batch.shape[0])
out_ok = jax.vmap(noisy_sample)(keys_b, xs_batch)
loop_ok = jnp.stack([noisy_sample(kb, xb) for kb, xb in zip(keys_b, xs_batch)])
print("vmap result matches calling one at a time:", bool(jnp.array_equal(out_ok, loop_ok)))
assert jnp.array_equal(out_ok, loop_ok)
assert len({np.asarray(r).tobytes() for r in out_ok}) == 4

# Two levels of vmap: split(key, (B, T)) gives a 2D key array directly
keys_bt = jax.random.split(jax.random.key(1), (4, 5))
out_bt = jax.vmap(jax.vmap(lambda kk: jax.random.normal(kk)))(keys_bt)
assert out_bt.shape == (4, 5) and out_bt[2, 3] == jax.random.normal(keys_bt[2, 3])

# ❌ Wrong: use the process index directly as the seed (process_index simulated with an argument here)
def process_key_bad(seed, process_index):
    return jax.random.key(process_index)          # the seed is ignored!
try:
    collide = bool(process_key_bad(0, 1) == process_key_bad(123, 1))
    print("No error, but process 1 gets the same key in the seed-0 and seed-123 experiments:", collide)
except Exception as e:
    print("❌", err_line(e))
assert collide

# ✅ Correct: the experiment seed determines the main key, then fold_in the process index
def process_key(seed, process_index):
    return jax.random.fold_in(jax.random.key(seed), process_index)
pkeys = [process_key(s, p) for s in (0, 123) for p in range(4)]
print("2 experiments × 4 processes = 8 keys, all distinct:",
      len({tuple(np.asarray(jax.random.key_data(x)).tolist()) for x in pkeys}) == 8)
assert len({tuple(np.asarray(jax.random.key_data(x)).tolist()) for x in pkeys}) == 8
assert bool(process_key(0, jax.process_index()) == jax.random.fold_in(jax.random.key(0), 0))   # index=0 in a single process

MP_WORKER = """
import os, sys, json
pid, nproc, port = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3]
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"   # 4 devices per process
import jax, jax.numpy as jnp, numpy as np
from jax.sharding import NamedSharding, PartitionSpec as P
jax.distributed.initialize(coordinator_address=f"localhost:{port}", num_processes=nproc, process_id=pid)
mesh = jax.make_mesh((jax.device_count(),), ("data",))
gen = jax.jit(lambda k: jax.random.normal(k, (8, 4)), out_shardings=NamedSharding(mesh, P("data")))
base = jax.random.key(42)
def local_rows(arr):      # shards this process can access: {starting row: data}
    return {int(s.index[0].start): np.asarray(s.data).tolist() for s in arr.addressable_shards}
out = {"pid": jax.process_index(), "nproc": jax.process_count(), "ndev": jax.device_count(),
       "shared": local_rows(gen(base)),                                          # ✅ same key
       "per_process": local_rows(gen(jax.random.fold_in(base, jax.process_index()))),   # ❌ a different key per process
       "local_perm": np.asarray(jax.random.permutation(jax.random.fold_in(base, jax.process_index()), 6)).tolist()}
print("JSON:" + json.dumps(out))
"""

def run_two_processes(code, nproc=2, timeout=300):
    s = socket.socket(); s.bind(("localhost", 0)); port = s.getsockname()[1]; s.close()
    procs = [subprocess.Popen([sys.executable, "-c", code, str(i), str(nproc), str(port)],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) for i in range(nproc)]
    results = {}
    for p in procs:
        o, e = p.communicate(timeout=timeout)
        assert p.returncode == 0, e[-3000:]
        r = json.loads([l for l in o.splitlines() if l.startswith("JSON:")][-1][5:])
        results[r["pid"]] = r
    return results

mp = run_two_processes(MP_WORKER)

def assemble(field):
    rows = {}
    for r in mp.values():
        rows.update({int(i): v for i, v in r[field].items()})
    return np.concatenate([np.array(rows[i]) for i in sorted(rows)])

ref42 = np.asarray(jax.random.normal(jax.random.key(42), (8, 4)))
ref_p0 = np.asarray(jax.random.normal(jax.random.fold_in(jax.random.key(42), 0), (8, 4)))
ref_p1 = np.asarray(jax.random.normal(jax.random.fold_in(jax.random.key(42), 1), (8, 4)))
g_shared, g_bad = assemble("shared"), assemble("per_process")

for pid, r in sorted(mp.items()):
    print(f"Process {pid}/{r['nproc']}: {r['ndev']} global devices, rows held by this process {sorted(int(i) for i in r['shared'])}, local shuffle {r['local_perm']}")
print("✅ Same key: assembled global array == single-process normal(key(42), (8,4))?", np.array_equal(g_shared, ref42))
print("❌ Different key per process: first 4 rows == first 4 rows of key A?", np.array_equal(g_bad[:4], ref_p0[:4]),
      "| last 4 rows == last 4 rows of key B?", np.array_equal(g_bad[4:], ref_p1[4:]),
      "| equal to any complete array?", any(np.array_equal(g_bad, r_) for r_ in (ref42, ref_p0, ref_p1)))

assert all(r["nproc"] == 2 and r["ndev"] == 8 for r in mp.values())
assert np.array_equal(g_shared, ref42)
assert np.array_equal(g_bad[:4], ref_p0[:4]) and np.array_equal(g_bad[4:], ref_p1[4:])
assert not any(np.array_equal(g_bad, r_) for r_ in (ref42, ref_p0, ref_p1))
assert mp[0]["local_perm"] != mp[1]["local_perm"]