# the ring algorithm: after step t, every device holds t+1 shards
mesh = Mesh(np.array(devs), ("x",))
AXIS = len(devs)

@partial(shard_map, mesh=mesh, in_specs=P("x", None), out_specs=P("x", None))
def ring_all_gather(shard):
    n = AXIS  # static mesh size from the closure
    idx = jax.lax.axis_index("x")
    out = jnp.zeros((n, *shard.shape[1:]), shard.dtype)
    out = out.at[idx].set(shard[0])
    recv = shard[0]
    def hop(t, carry):
        out, recv = carry
        recv = jax.lax.ppermute(recv, "x", [(i, (i + 1) % n) for i in range(n)])
        src = (idx - t - 1) % n
        return out.at[src].set(recv), recv
    out, _ = jax.lax.fori_loop(0, n - 1, hop, (out, recv))
    return out[None]

x = jax.random.normal(jax.random.key(0), (AXIS, 4, 128))
got = ring_all_gather(x).reshape(AXIS, AXIS, 4, 128)
want = jnp.broadcast_to(x[None], (AXIS, AXIS, 4, 128))
check("ring all-gather == replicated truth", got, want, tol=0)
print(f"validated over {AXIS} devices: every device reconstructed all {AXIS} shards")

# LAB 3.1's summarize/combine, verbatim: transport changes, algebra does not
def summarize(s_block, v_block):
    m = jnp.max(s_block, axis=-1)
    p = jnp.exp(s_block - m[..., None])
    return (m, jnp.sum(p, axis=-1), p @ v_block)

def combine(a, b):
    m1, l1, acc1 = a
    m2, l2, acc2 = b
    m = jnp.maximum(m1, m2)
    a1, a2 = jnp.exp(m1 - m), jnp.exp(m2 - m)
    return (m, l1 * a1 + l2 * a2, acc1 * a1[..., None] + acc2 * a2[..., None])

mesh = Mesh(np.array(devs), ("x",))
AXIS = len(devs)

@partial(shard_map, mesh=mesh,
         in_specs=(P("x", None), P("x", None), P("x", None)), out_specs=P("x", None))
def ring_attention(q, k, v):
    n = AXIS  # static mesh size from the closure
    perm = [(i, (i + 1) % n) for i in range(n)]
    state = summarize(q @ k.T, v)                    # my resident KV shard
    def hop(t, carry):
        state, k, v = carry
        k = jax.lax.ppermute(k, "x", perm)           # the shard rides the ring...
        v = jax.lax.ppermute(v, "x", perm)
        return combine(state, summarize(q @ k.T, v)), k, v   # ...while we fold it in
    state, _, _ = jax.lax.fori_loop(0, n - 1, hop, (state, k, v))
    m, l, acc = state
    return acc / l[..., None]

SQ, SKV, D = 16 * AXIS, 32 * AXIS, 64
q = jax.random.normal(jax.random.key(0), (SQ, D))
k = jax.random.normal(jax.random.key(1), (SKV, D))
v = jax.random.normal(jax.random.key(2), (SKV, D))

got = ring_attention(q, k, v)
want = jax.nn.softmax(q @ k.T, axis=-1) @ v
check("ring attention == full attention", got, want, tol=1e-5)
print(f"distributed over {AXIS} devices, no device ever held more than 1/{AXIS} of KV")