Upload iris/test_all.py

iris/test_all.py

@@ -0,0 +1,221 @@
+"""Comprehensive test suite for IRIS. 17 tests, all verified passing."""
+
+import torch
+import torch.nn.functional as F
+import time, sys, traceback, math
+
+def test_module(name, fn):
+    print(f"\n{'='*60}\nTEST: {name}\n{'='*60}")
+    try:
+        fn()
+        print(f"  PASSED")
+        return True
+    except Exception as e:
+        print(f"  FAILED: {e}")
+        traceback.print_exc()
+        return False
+
+def test_spectral_conv():
+    from iris.pde_ssm import SpectralConv2d
+    for H, W in [(4,4),(8,8),(16,16)]:
+        conv = SpectralConv2d(channels=32, modes_h=H//2, modes_w=W//2)
+        x = torch.randn(2, 32, H, W)
+        out = conv(x)
+        assert out.shape == x.shape
+        out.sum().backward()
+        assert conv.weight_pos.grad is not None and conv.weight_pos.grad.norm() > 0
+        assert not torch.isnan(out).any()
+        print(f"  SpectralConv2d({H}x{W}): OK")
+
+def test_token_differential():
+    from iris.pde_ssm import TokenDifferential
+    td = TokenDifferential(32)
+    x = torch.randn(2, 32, 4, 4)
+    assert torch.allclose(td(x), x, atol=1e-6)
+    td.alpha.data.fill_(1.0)
+    assert not torch.allclose(td(x), x)
+    print("  TokenDiff: identity at init, non-identity with alpha=1")
+
+def test_pde_ssm_block():
+    from iris.pde_ssm import PDESSMBlock
+    for s in [4, 8]:
+        b = PDESSMBlock(dim=64, spatial_size=s)
+        x = torch.randn(2, s*s, 64)
+        out = b(x, s, s)
+        assert out.shape == x.shape
+        out.sum().backward()
+        for n, p in b.named_parameters():
+            if p.requires_grad:
+                assert p.grad is not None and not torch.isnan(p.grad).any()
+        print(f"  PDESSMBlock(s={s}): OK")
+
+def test_cross_attention():
+    from iris.blocks import MultiQueryCrossAttention
+    a = MultiQueryCrossAttention(dim=64, num_heads=4)
+    out = a(torch.randn(2,16,64), torch.randn(2,32,64))
+    assert out.shape == (2,16,64)
+    out.sum().backward()
+    assert a.k_proj.weight.numel() < a.q_proj.weight.numel()
+    print(f"  CrossAttn MQA: OK, K/Q ratio = {a.q_proj.weight.numel()//a.k_proj.weight.numel()}x")
+
+def test_self_attention():
+    from iris.blocks import MultiQuerySelfAttention
+    a = MultiQuerySelfAttention(dim=64, num_heads=4)
+    out = a(torch.randn(2,16,64), 4, 4)
+    assert out.shape == (2,16,64)
+    out.sum().backward()
+    print("  SelfAttn+2D RoPE: OK")
+
+def test_rope_2d():
+    from iris.blocks import RotaryEmbedding2D
+    rope = RotaryEmbedding2D(dim=16)
+    x = torch.randn(2, 4, 16, 16)
+    out = rope(x, 4, 4)
+    assert out.shape == x.shape
+    assert abs(x.norm(dim=-1).mean() - out.norm(dim=-1).mean()) / x.norm(dim=-1).mean() < 0.1
+    print("  2D RoPE: norm preserved")
+
+def test_uib_ffn():
+    from iris.blocks import UIBFFN
+    f = UIBFFN(dim=64, expansion=2)
+    out = f(torch.randn(2,16,64), 4, 4)
+    assert out.shape == (2,16,64)
+    out.sum().backward()
+    assert f.dw_conv.groups == 128
+    print("  UIB-FFN: OK")
+
+def test_timestep_embedding():
+    from iris.blocks import TimestepEmbedding
+    te = TimestepEmbedding(dim=64)
+    out = te(torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0]))
+    assert out.shape == (5, 64) and not torch.isnan(out).any()
+    print("  TimestepEmbed: OK")
+
+def test_patchify_unpatchify():
+    from iris.model import Patchify, Unpatchify
+    for ps in [2, 4]:
+        dim = 128 if ps==2 else 512
+        p, u = Patchify(32, dim, ps), Unpatchify(32, dim, ps)
+        z = torch.randn(2, 32, 16, 16)
+        tok, H, W = p(z)
+        assert tok.shape == (2, (16//ps)**2, dim)
+        assert u(tok, H, W).shape == z.shape
+        print(f"  Patchify(ps={ps}): OK")
+
+def test_tiny_decoder():
+    from iris.model import TinyDecoder
+    d = TinyDecoder(32, 3)
+    img = d(torch.randn(2, 32, 16, 16))
+    assert img.shape == (2, 3, 512, 512)
+    n = sum(p.numel() for p in d.parameters())
+    assert n < 2_000_000
+    print(f"  TinyDecoder: {n:,} params, output {img.shape}")
+
+def test_iris_forward():
+    from iris.model import IRIS
+    m = IRIS(latent_channels=32, dim=128, patch_size=4, num_blocks=4, num_heads=4, max_iterations=4, gradient_checkpointing=False)
+    z = torch.randn(2, 32, 16, 16)
+    t = torch.tensor([0.3, 0.7])
+    ctx = torch.randn(2, 8, 128)
+    v = m(z, t, ctx, num_iterations=2)
+    assert v.shape == z.shape and not torch.isnan(v).any()
+    v.sum().backward()
+    opt = torch.optim.SGD(m.parameters(), lr=0.01)
+    opt.step(); opt.zero_grad(set_to_none=True)
+    m(z, t, ctx, 2).sum().backward()
+    core_p = [(n,p) for n,p in m.named_parameters() if p.requires_grad and "tiny_decoder" not in n]
+    assert all(p.grad is not None and p.grad.norm()>1e-10 for _,p in core_p)
+    print(f"  Forward OK, all {len(core_p)} core params have grad, total={m.count_params()['total']:,}")
+
+def test_flow_matching_loss():
+    from iris.model import IRIS
+    from iris.flow_matching import flow_matching_loss
+    m = IRIS(latent_channels=32, dim=128, patch_size=4, num_blocks=4, num_heads=4, max_iterations=4, gradient_checkpointing=False)
+    l = flow_matching_loss(m, torch.randn(4,32,16,16)*2.5, torch.randn(4,8,128), num_iterations=2)
+    assert l["loss"].requires_grad and not torch.isnan(l["loss"]) and l["loss"].item() > 0
+    l["loss"].backward()
+    print(f"  flow_loss={l['flow_loss'].item():.4f}")
+
+def test_euler_sampling():
+    from iris.model import IRIS
+    from iris.flow_matching import euler_sample
+    m = IRIS(latent_channels=32, dim=128, patch_size=4, num_blocks=4, num_heads=4, max_iterations=4, gradient_checkpointing=False)
+    m.eval()
+    with torch.no_grad():
+        z = euler_sample(m, torch.randn(1,32,16,16), torch.randn(1,8,128), num_steps=5, num_iterations=2)
+    assert z.shape == (1,32,16,16) and not torch.isnan(z).any()
+    img = m.decode_latent(z)
+    assert img.shape == (1,3,512,512)
+    print(f"  Euler sampling OK, decoded {img.shape}")
+
+def test_gradient_checkpointing():
+    from iris.model import IRIS
+    from iris.flow_matching import flow_matching_loss
+    m1 = IRIS(latent_channels=32, dim=64, patch_size=4, num_blocks=3, num_heads=4, max_iterations=4, gradient_checkpointing=False)
+    m2 = IRIS(latent_channels=32, dim=64, patch_size=4, num_blocks=3, num_heads=4, max_iterations=4, gradient_checkpointing=True)
+    m2.load_state_dict(m1.state_dict())
+    torch.manual_seed(42)
+    z, ctx = torch.randn(2,32,16,16)*2.5, torch.randn(2,4,64)
+    torch.manual_seed(123)
+    l1 = flow_matching_loss(m1, z, ctx, num_iterations=3); l1["loss"].backward()
+    torch.manual_seed(123)
+    l2 = flow_matching_loss(m2, z, ctx, num_iterations=3); l2["loss"].backward()
+    diff = abs(l1["loss"].item() - l2["loss"].item())
+    maxg = max((p1.grad-p2.grad).abs().max().item() for (n1,p1),(n2,p2) in zip(m1.named_parameters(),m2.named_parameters()) if p1.grad is not None and p2.grad is not None)
+    assert diff < 1e-6 and maxg < 1e-4
+    print(f"  Checkpointing: loss diff={diff:.8f}, max grad diff={maxg:.8f}")
+
+def test_weight_sharing():
+    from iris.core import RefinementCore
+    c = RefinementCore(dim=64, num_blocks=3, num_heads=4, spatial_size=4, max_iterations=4, gradient_checkpointing=False)
+    x, ctx, t = torch.randn(1,16,64), torch.randn(1,4,64), torch.tensor([0.5])
+    o2 = c(x, ctx, t, 4, 4, num_iterations=2)
+    o4 = c(x, ctx, t, 4, 4, num_iterations=4)
+    assert not torch.allclose(o2, o4, atol=1e-3)
+    print(f"  Weight sharing: {sum(p.numel() for p in c.parameters()):,} params (constant)")
+
+def test_zero_init():
+    from iris.model import IRIS
+    m = IRIS(latent_channels=32, dim=64, patch_size=4, num_blocks=3, num_heads=4, gradient_checkpointing=False)
+    assert (m.unpatchify.proj.weight==0).all() and (m.unpatchify.proj.bias==0).all()
+    with torch.no_grad():
+        out = m(torch.randn(1,32,16,16), torch.tensor([0.5]), torch.randn(1,4,64), 1)
+    assert out.norm().item() < 1.0
+    print(f"  Zero-init: output norm={out.norm().item():.6f}")
+
+def test_training_stability():
+    from iris.model import IRIS
+    from iris.flow_matching import flow_matching_loss
+    m = IRIS(latent_channels=32, dim=128, patch_size=4, num_blocks=4, num_heads=4, max_iterations=4, gradient_checkpointing=False)
+    opt = torch.optim.AdamW(m.parameters(), lr=1e-3, weight_decay=0.01)
+    torch.manual_seed(0)
+    z, ctx = torch.randn(8,32,16,16)*2.5, torch.randn(8,8,128)
+    losses = []
+    m.train()
+    for s in range(100):
+        l = flow_matching_loss(m, z, ctx, num_iterations=2)
+        opt.zero_grad(set_to_none=True)
+        l["loss"].backward()
+        torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
+        opt.step()
+        losses.append(l["loss"].item())
+        if (s+1) % 25 == 0: print(f"  Step {s+1}: loss={losses[-1]:.4f}")
+    f10, l10 = sum(losses[:10])/10, sum(losses[-10:])/10
+    print(f"  Loss: {f10:.4f} -> {l10:.4f} ({(1-l10/f10)*100:.1f}%)")
+    assert l10 < f10 and not any(math.isnan(l) or math.isinf(l) for l in losses)
+
+if __name__ == "__main__":
+    tests = [
+        ("SpectralConv2d", test_spectral_conv), ("TokenDifferential", test_token_differential),
+        ("PDESSMBlock", test_pde_ssm_block), ("CrossAttention (MQA)", test_cross_attention),
+        ("SelfAttention (MQA+2D RoPE)", test_self_attention), ("2D RoPE", test_rope_2d),
+        ("UIB-FFN", test_uib_ffn), ("TimestepEmbedding", test_timestep_embedding),
+        ("Patchify/Unpatchify", test_patchify_unpatchify), ("TinyDecoder", test_tiny_decoder),
+        ("IRIS Forward", test_iris_forward), ("Flow Matching Loss", test_flow_matching_loss),
+        ("Euler Sampling", test_euler_sampling), ("Gradient Checkpointing", test_gradient_checkpointing),
+        ("Weight Sharing", test_weight_sharing), ("Zero-Init Output", test_zero_init),
+        ("Training Stability (100 steps)", test_training_stability),
+    ]
+    passed = sum(1 for n,f in tests if test_module(n,f))
+    print(f"\n{'='*60}\nRESULTS: {passed}/{len(tests)} passed\n{'='*60}")
+    if passed < len(tests): sys.exit(1)