"""DiT block tests against torch. These layers all have torch counterparts, so unlike the submanifold conv they are checked against upstream's real semantics. The upstream forward bodies are reproduced verbatim from pixal3d/modules/{attention/modules.py,transformer/modulated.py}. """ import sys from pathlib import Path import mlx.core as mx import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from trellis_sparse_mlx.dit import ( # noqa: E402 DiTAttention, ModulatedTransformerCrossBlock, MultiHeadRMSNorm, TimestepEmbedder, apply_rope, rope_phases_from_coords, ) def test_rms_norm_matches_upstream(dim=32, heads=4, n=17): """Upstream writes it as F.normalize(x)*gamma*sqrt(dim); we implement RMS directly.""" rng = np.random.default_rng(0) x = rng.standard_normal((n, heads, dim)).astype(np.float32) g = rng.standard_normal((heads, dim)).astype(np.float32) m = MultiHeadRMSNorm(dim, heads) m.gamma = mx.array(g) got = np.asarray(m(mx.array(x))) tx = torch.tensor(x) want = (F.normalize(tx.float(), dim=-1) * torch.tensor(g) * dim**0.5).numpy() err = np.abs(got - want).max() assert err < 2e-4, f"rms norm err {err:.3g}" return err def test_rms_norm_is_scale_invariant(): """RMS norm must remove input magnitude — catches a plain scale-by-gamma stub.""" m = MultiHeadRMSNorm(8, 2) x = np.random.default_rng(1).standard_normal((5, 2, 8)).astype(np.float32) a = np.asarray(m(mx.array(x))) b = np.asarray(m(mx.array(x * 37.0))) err = np.abs(a - b).max() assert err < 1e-3, f"not scale-invariant: {err:.3g}" return err def test_rope_is_a_rotation(): """RoPE must preserve the norm of each rotary pair.""" rng = np.random.default_rng(2) q = rng.standard_normal((2, 6, 4, 16)).astype(np.float32) ph = rng.uniform(0, 2 * np.pi, (2, 6, 4, 8)).astype(np.float32) rq, _ = apply_rope(mx.array(q), mx.array(q), mx.array(ph)) rq = np.asarray(rq) n0 = np.sqrt(q[..., 0::2] ** 2 + q[..., 1::2] ** 2) n1 = np.sqrt(rq[..., 0::2] ** 2 + rq[..., 1::2] ** 2) err = np.abs(n0 - n1).max() assert err < 1e-4, f"rope changed pair norms by {err:.3g}" # and it must actually rotate assert np.abs(rq - q).max() > 1e-3, "rope was a no-op" return err def test_attention_matches_torch(ch=64, heads=8, n=13, b=2): rng = np.random.default_rng(3) x = rng.standard_normal((b, n, ch)).astype(np.float32) wq = rng.standard_normal((ch * 3, ch)).astype(np.float32) * 0.05 bq = rng.standard_normal((ch * 3,)).astype(np.float32) * 0.05 wo = rng.standard_normal((ch, ch)).astype(np.float32) * 0.05 bo = rng.standard_normal((ch,)).astype(np.float32) * 0.05 a = DiTAttention(ch, heads) a.to_qkv.weight, a.to_qkv.bias = mx.array(wq), mx.array(bq) a.to_out.weight, a.to_out.bias = mx.array(wo), mx.array(bo) got = np.asarray(a(mx.array(x))) d = ch // heads tx = torch.tensor(x) qkv = F.linear(tx, torch.tensor(wq), torch.tensor(bq)).reshape(b, n, 3, heads, d) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] o = F.scaled_dot_product_attention( q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3) ) o = o.permute(0, 2, 1, 3).reshape(b, n, ch) want = F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy() err = np.abs(got - want).max() assert err < 2e-4, f"attention err {err:.3g}" return err def test_modulation_gates_actually_gate(ch=32, heads=4, n=7, b=2): """gate=0 must make the block an identity on the modulated paths.""" blk = ModulatedTransformerCrossBlock(ch, ch, heads, share_mod=True) x = mx.array(np.random.default_rng(4).standard_normal((b, n, ch)).astype(np.float32)) ctx = mx.array(np.zeros((b, 5, ch), dtype=np.float32)) # zero everything, then force the cross-attn output to zero via zero out-proj blk.modulation = mx.zeros((6 * ch,)) blk.cross_attn.to_out.weight = mx.zeros_like(blk.cross_attn.to_out.weight) blk.cross_attn.to_out.bias = mx.zeros_like(blk.cross_attn.to_out.bias) mod = mx.zeros((b, 6 * ch)) out = np.asarray(blk(x, mod, ctx)) err = np.abs(out - np.asarray(x)).max() assert err < 1e-4, f"gates did not zero the residual branches: {err:.3g}" return err def test_norm_affine_asymmetry(): """norm1/norm3 non-affine, norm2 affine — silent if swapped.""" blk = ModulatedTransformerCrossBlock(16, 16, 2, share_mod=True) assert not blk.norm1.affine, "norm1 should be non-affine" assert blk.norm2.affine, "norm2 SHOULD be affine" assert not blk.norm3.affine, "norm3 should be non-affine" return 0.0 def test_timestep_embedder_shape(): te = TimestepEmbedder(64, 32) out = te(mx.array(np.array([0.0, 0.5, 1.0], dtype=np.float32))) assert out.shape == (3, 64), out.shape assert np.isfinite(np.asarray(out)).all() return 0.0 def test_rope_and_rms_do_not_commute(): """Regression: upstream applies qk RMS norm BEFORE RoPE, and the two do not commute. Reversed, a single block still correlates 0.9998 with upstream — which compounds to 0.84 over 30 blocks. This asserts the operations genuinely differ by order, so the ordering in DiTAttention is load-bearing and not cosmetic. """ rng = np.random.default_rng(11) h, d, n = 4, 16, 9 q = mx.array(rng.standard_normal((1, n, h, d)).astype(np.float32)) ph = mx.array(np.exp(1j * rng.uniform(0, 2 * np.pi, (n, d // 2))).astype(np.complex64)) rms = MultiHeadRMSNorm(d, h) rms.gamma = mx.array(rng.standard_normal((h, d)).astype(np.float32)) norm_then_rope, _ = apply_rope(rms(q), rms(q), ph) rope_first, _ = apply_rope(q, q, ph) rope_then_norm = rms(rope_first) gap = np.abs(np.asarray(norm_then_rope) - np.asarray(rope_then_norm)).max() assert gap > 1e-3, f"orders agree ({gap:.3g}) — test cannot catch a regression" return gap def test_complex_rope_matches_torch_complex_multiply(): """rope_phases ships complex; the rotation must be a complex multiply, not cos(phase).""" rng = np.random.default_rng(12) b, n, h, d = 1, 7, 3, 8 q = rng.standard_normal((b, n, h, d)).astype(np.float32) ph = np.exp(1j * rng.uniform(0, 2 * np.pi, (n, d // 2))).astype(np.complex64) got, _ = apply_rope(mx.array(q), mx.array(q), mx.array(ph)) tq = torch.view_as_complex(torch.tensor(q).reshape(b, n, h, d // 2, 2).contiguous()) want = torch.view_as_real(tq * torch.tensor(ph)[None, :, None, :]).reshape(b, n, h, d) err = np.abs(np.asarray(got) - want.numpy()).max() assert err < 1e-5, f"complex rope err {err:.3g}" return err def test_derived_rope_matches_shipped_tensor(): """The derivation must reproduce Pixal3D's own precomputed rope_phases. ss_flow ships `rope_phases` for a 16^3 grid, so there is a ground-truth tensor to check against — including the detail that 3*21=63 frequencies fall one short of head_dim//2=64 and upstream right-pads with 1+0j. """ ck = Path( "/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/" "ss_flow_img_dit_1_3B_64_bf16.safetensors" ) if not ck.exists(): print(" (skipped: checkpoint not present)") return 0.0 ref = np.asarray(mx.load(str(ck))["rope_phases"]) res = 16 zz, yy, xx = np.meshgrid(*[np.arange(res)] * 3, indexing="ij") coords = np.stack([zz, yy, xx], -1).reshape(-1, 3).astype(np.int32) got = np.asarray(rope_phases_from_coords(mx.array(coords), head_dim=128)) err = np.abs(got - ref).max() assert err < 1e-5, f"derived phases differ from shipped by {err:.3g}" return err if __name__ == "__main__": tests = [ ("rms norm vs upstream", test_rms_norm_matches_upstream), ("rms scale invariance", test_rms_norm_is_scale_invariant), ("rope is a rotation", test_rope_is_a_rotation), ("dit attention vs torch", test_attention_matches_torch), ("modulation gates", test_modulation_gates_actually_gate), ("norm affine asymmetry", test_norm_affine_asymmetry), ("timestep embedder", test_timestep_embedder_shape), ("rope/rms don't commute", test_rope_and_rms_do_not_commute), ("complex rope vs torch", test_complex_rope_matches_torch_complex_multiply), ("derived rope vs shipped", test_derived_rope_matches_shipped_tensor), ] failed = 0 for name, fn in tests: try: err = fn() print(f" PASS {name:26s} (max err {err:.2e})") except AssertionError as e: print(f" FAIL {name:26s} {e}") failed += 1 except Exception as e: # noqa: BLE001 print(f" ERROR {name:26s} {type(e).__name__}: {e}") failed += 1 print(f"\n{len(tests)-failed}/{len(tests)} passed") sys.exit(1 if failed else 0)