run_inference.py

"""
Simple inference script - just run models with random data.
No profiling, no timing - you can add that yourself.
"""

import torch
import torch_xla
import torch_xla.debug.profiler as xp
import argparse
from models import (
    EBM, DiffusionWrapper,
    SudokuEBM, SudokuTransformerEBM, SudokuDenoise, SudokuLatentEBM,
    GraphEBM, GraphReverse, 
    GNNConvEBM, GNNDiffusionWrapper, GNNConvDiffusionWrapper,
    GNNConv1DEBMV2, GNNConv1DV2DiffusionWrapper, GNNConv1DReverse
)


def run_mlp(batch_size=64):
    """MLP model for continuous tasks (addition, inverse, lowrank)"""
    device = torch_xla.device()
    
    # 20x20 matrix = 400 dim
    inp_dim = 400
    out_dim = 400
    
    model = EBM(inp_dim, out_dim, is_ebm=True).to(device)
    model = DiffusionWrapper(model).to(device)
    model.eval()
    
    # Random inputs (DiffusionWrapper needs gradients)
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    # Forward pass
    xp.start_trace("./traces/mlp_forward")
    output = model(inp, out, t)
    xp.stop_trace()
    
    print(f"MLP: input={inp.shape}, output={output.shape}")
    return output


def run_mlp_reverse(batch_size=64):
    """MLP reverse (denoising) model"""
    device = torch_xla.device()
    
    inp_dim = 400
    out_dim = 400
    
    model = EBM(inp_dim, out_dim, is_ebm=False).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    # No wrapper, can use no_grad
    with torch.no_grad():
        output = model(inp, out, t)
    
    print(f"MLP-Reverse: input={inp.shape}, output={output.shape}")
    return output


def run_sudoku(batch_size=64):
    """Sudoku model"""
    device = torch_xla.device()
    
    # 9x9 grid with 9 channels = 729
    inp_dim = 729
    out_dim = 729
    
    model = SudokuEBM(inp_dim, out_dim).to(device)
    model = DiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"Sudoku: input={inp.shape}, output={output.shape}")
    return output


def run_sudoku_latent(batch_size=64):
    """Sudoku latent model"""
    device = torch_xla.device()
    
    inp_dim = 729
    out_dim = 243  # Latent representation: 9x9x3 = 243
    
    model = SudokuLatentEBM(inp_dim, out_dim).to(device)
    model = DiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"Sudoku-Latent: input={inp.shape}, output={output.shape}")
    return output


def run_sudoku_transformer(batch_size=64):
    """Sudoku transformer model"""
    device = torch_xla.device()
    
    inp_dim = 729
    out_dim = 729
    
    model = SudokuTransformerEBM(inp_dim, out_dim).to(device)
    model = DiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"Sudoku-Transformer: input={inp.shape}, output={output.shape}")
    return output


def run_sudoku_reverse(batch_size=64):
    """Sudoku reverse (denoising) model"""
    device = torch_xla.device()
    
    inp_dim = 729
    out_dim = 729
    
    model = SudokuDenoise(inp_dim, out_dim).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, inp_dim, device=device)
    out = torch.randn(batch_size, out_dim, device=device)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    # No wrapper, can use no_grad
    with torch.no_grad():
        output = model(inp, out, t)
    
    print(f"Sudoku-Reverse: input={inp.shape}, output={output.shape}")
    return output


def run_gnn(batch_size=64):
    """GNN model for graph connectivity"""
    device = torch_xla.device()
    
    # 12x12 graph with 4 features per edge
    n_nodes = 12
    inp_dim = 4
    out_dim = 1
    
    model = GraphEBM(inp_dim, out_dim).to(device)
    model = GNNDiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, n_nodes, n_nodes, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"GNN: input={inp.shape}, output={output.shape}")
    return output


def run_gnn_reverse(batch_size=64):
    """GNN reverse model"""
    device = torch_xla.device()
    
    n_nodes = 12
    inp_dim = 4
    out_dim = 1
    
    model = GraphReverse(inp_dim, out_dim).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, n_nodes, n_nodes, device=device)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    # No wrapper, can use no_grad
    with torch.no_grad():
        output = model(inp, out, t)
    
    print(f"GNN-Reverse: input={inp.shape}, output={output.shape}")
    return output


def run_gnn_conv(batch_size=64):
    """GNN conv model"""
    device = torch_xla.device()
    
    n_nodes = 12
    inp_dim = 4
    out_dim = 1
    seq_len = 16
    
    model = GNNConvEBM(inp_dim, out_dim, use_1d=False).to(device)
    model = GNNConvDiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, seq_len, n_nodes, n_nodes, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"GNN-Conv: input={inp.shape}, output={output.shape}")
    return output


def run_gnn_conv_1d(batch_size=64):
    """GNN conv 1D model"""
    device = torch_xla.device()
    
    n_nodes = 12
    inp_dim = 2
    out_dim = 1
    seq_len = 16
    
    model = GNNConvEBM(inp_dim, out_dim, use_1d=True).to(device)
    model = GNNConvDiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, seq_len, n_nodes, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"GNN-Conv-1D: input={inp.shape}, output={output.shape}")
    return output


def run_gnn_conv_1d_v2(batch_size=64):
    """GNN conv 1D v2 model (for planning)"""
    device = torch_xla.device()
    
    n_nodes = 20
    inp_dim = 4
    out_dim = 1
    seq_len = 8
    
    model = GNNConv1DEBMV2(inp_dim, out_dim).to(device)
    model = GNNConv1DV2DiffusionWrapper(model).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, seq_len, n_nodes, out_dim, device=device, requires_grad=True)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    output = model(inp, out, t)
    
    print(f"GNN-Conv-1D-V2: input={inp.shape}, output={output.shape}")
    return output


def run_gnn_conv_1d_v2_reverse(batch_size=64):
    """GNN conv 1D v2 reverse model"""
    device = torch_xla.device()
    
    n_nodes = 20
    inp_dim = 4
    out_dim = 1
    seq_len = 8
    
    model = GNNConv1DReverse(inp_dim, out_dim).to(device)
    model.eval()
    
    inp = torch.randn(batch_size, n_nodes, n_nodes, inp_dim, device=device)
    out = torch.randn(batch_size, seq_len, n_nodes, out_dim, device=device)
    t = torch.zeros(batch_size, dtype=torch.long, device=device)
    
    # No wrapper, can use no_grad
    with torch.no_grad():
        output = model(inp, out, t)
    
    print(f"GNN-Conv-1D-V2-Reverse: input={inp.shape}, output={output.shape}")
    return output


def main():
    parser = argparse.ArgumentParser(description='Run IRED model inference with random data')
    parser.add_argument('--model', type=str, default='mlp',
                       choices=['mlp', 'mlp-reverse', 'sudoku', 'sudoku-latent', 
                               'sudoku-transformer', 'sudoku-reverse', 'gnn', 'gnn-reverse',
                               'gnn-conv', 'gnn-conv-1d', 'gnn-conv-1d-v2', 
                               'gnn-conv-1d-v2-reverse', 'all'],
                       help='Which model to run')
    parser.add_argument('--batch-size', type=int, default=64, help='Batch size')
    
    args = parser.parse_args()
    
    device = torch_xla.device()
    print(f"Device: {device}")
    print(f"Batch size: {args.batch_size}\n")
    
    if args.model == 'all' or args.model == 'mlp':
        run_mlp(args.batch_size)
    
    if args.model == 'all' or args.model == 'mlp-reverse':
        run_mlp_reverse(args.batch_size)
    
    if args.model == 'all' or args.model == 'sudoku':
        run_sudoku(args.batch_size)
    
    if args.model == 'all' or args.model == 'sudoku-latent':
        run_sudoku_latent(args.batch_size)
    
    if args.model == 'all' or args.model == 'sudoku-transformer':
        run_sudoku_transformer(args.batch_size)
    
    if args.model == 'all' or args.model == 'sudoku-reverse':
        run_sudoku_reverse(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn':
        run_gnn(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn-reverse':
        run_gnn_reverse(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn-conv':
        run_gnn_conv(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn-conv-1d':
        run_gnn_conv_1d(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn-conv-1d-v2':
        run_gnn_conv_1d_v2(args.batch_size)
    
    if args.model == 'all' or args.model == 'gnn-conv-1d-v2-reverse':
        run_gnn_conv_1d_v2_reverse(args.batch_size)
    
    print("\nDone!")


if __name__ == '__main__':
    server = xp.start_server(9012)
    main()