"""Synthetic checks accompanying the SwiGLU methods note; not a model benchmark.

Run with Python and NumPy: python examples/swiglu_check.py
"""
import numpy as np

tiny = np.array([1e-30], dtype=np.float32)
assert np.max(np.abs(tiny)) > 0
assert np.square(tiny)[0] == 0

rng = np.random.default_rng(7)
up = rng.normal(0, 0.2, (4, 3)).astype(np.float32)
gate = rng.normal(0, 0.2, (4, 3)).astype(np.float32)
down = rng.normal(0, 0.2, (2, 4)).astype(np.float32)
up[1, :] = 1e-31
gate[1, :] = 1e-31
down[:, 1] = 1e-31
# A small row in only one matrix does not pass the conservative joint rule.
up[2, :] = 1e-31
maximum = np.maximum.reduce([
    np.abs(up).max(axis=1),
    np.abs(gate).max(axis=1),
    np.abs(down).max(axis=0),
])
remove = maximum < 1e-30
assert remove.tolist() == [False, True, False, False]
keep = ~remove

def forward(x, u, g, d):
    z = x @ u.T
    silu = z / (1.0 + np.exp(-z))
    return (silu * (x @ g.T)) @ d.T

x = rng.normal(0, 0.3, (8, 3)).astype(np.float32)
original = forward(x, up, gate, down)
pruned = forward(x, up[keep, :], gate[keep, :], down[:, keep])
np.testing.assert_allclose(original, pruned, rtol=1e-6, atol=1e-7)
print('Synthetic indexing, underflow, and output checks passed.')
print('Maximum absolute output difference:', float(np.abs(original-pruned).max()))