import torch
from torch.utils.data import DataLoader, Dataset
from torch._higher_order_ops.cudagraph_conditional_nodes import (
    ControlFlowOpWarmupDispatchMode,
    CUDAGraphCaptureControlFlowOpDispatchMode,
)

class RandomDataset(Dataset):
    def __init__(self, n_samples=256, in_dim=16, out_dim=4):
        self.n_samples = n_samples
        self.in_dim = in_dim
        self.out_dim = out_dim

    def __len__(self):
        return self.n_samples

    def __getitem__(self, _):
        return torch.randn(self.in_dim), torch.randn(self.out_dim)

model = torch.nn.Linear(16, 4).cuda()
criterion = torch.nn.MSELoss()
data_loader = DataLoader(RandomDataset(), batch_size=32, shuffle=False,
                         pin_memory=True, num_workers=1)


lr = torch.tensor(0.1, device="cuda")
loss_scale = torch.tensor(2.0**8, device="cuda")

def train(X, y):
    with torch.autograd.grad_mode.set_multithreading_enabled(False):
        global loss_scale
        model.zero_grad(set_to_none=True)
        y_pred = model(X)
        loss = criterion(y_pred, y)
        (loss * loss_scale).backward()
        all_gradients_finite = torch.stack([
            torch.all(param.grad.isfinite()) for param in model.parameters() if param.grad is not None
        ]).all()
        with torch.no_grad():
            def true_fn():
                global lr, loss_scale
                lr.mul_(0.99999)
                inv_scale = loss_scale.reciprocal()
                for param in model.parameters():
                    if param.grad is not None:
                        param.sub_(lr * param.grad * inv_scale)
                loss_scale.mul_(2.0)

            def false_fn():
                loss_scale.mul_(0.5)

            _ = torch.cond(all_gradients_finite, true_fn, false_fn)
        return loss

for i, (X, y) in enumerate(data_loader):
    if i == 0:
        with ControlFlowOpWarmupDispatchMode():
            loss = train(X.cuda(), y.cuda())
        print(loss.item())
        del loss
        continue
    if i == 1:
        X_static = torch.empty_like(X, device="cuda")
        y_static = torch.empty_like(y, device="cuda")
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g, capture_error_mode="thread_local"), CUDAGraphCaptureControlFlowOpDispatchMode():
            loss_static = train(X_static, y_static)
    if i == 2:
        torch.cuda.cudart().cudaProfilerStart()
    X_static.copy_(X, non_blocking=True)
    y_static.copy_(y, non_blocking=True)
    g.replay()
    print(loss_static.item())