simple_inference.py

#!/usr/bin/env python3
"""
Simple inference - just run the model once with random data on TPU
"""
import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

import torch
import numpy as np
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.profiler as xp
from diffusion_lib.denoising_diffusion_pytorch_1d import GaussianDiffusion1D
from models import EBM, DiffusionWrapper


def main():
    # Settings (matching your training)
    xp.start_server(9012)
    batch_size = 2048
    rank = 20
    inp_dim = 2 * rank * rank  # 800
    out_dim = rank * rank       # 400
    
    # Use TPU device
    device = xm.xla_device()
    print(f"Device: {device}")
    
    # Create model
    model = EBM(inp_dim=inp_dim, out_dim=out_dim)
    model = DiffusionWrapper(model)
    
    diffusion = GaussianDiffusion1D(
        model,
        seq_length=32,
        objective='pred_noise',
        timesteps=3,
        sampling_timesteps=3,
        supervise_energy_landscape=True,
        use_innerloop_opt=True,
        continuous=True,
        show_inference_tqdm=False,
    )
    
    diffusion = diffusion.to(device)
    diffusion.eval()
    
    # Generate random input data (don't care about the actual values)
    inp_batch = torch.randn(batch_size, inp_dim, device=device)
    
    # Run inference once
    print(f"Running inference with batch_size={batch_size}...")
    torch.set_grad_enabled(False)

    xp.start_trace("./traces/addition")
    samples = diffusion.sample(
        x=inp_batch,
        label=None,
        mask=None,
        batch_size=batch_size
    )
    
    # Synchronize TPU operations
    torch_xla.sync()
    
    print(f"Done! Output shape: {samples.shape}")
    print(f"✅ Inference completed on TPU!")
    xp.stop_trace()

if __name__ == "__main__":
    main()