import numpy as np
import torch
import torch.utils.checkpoint
import torch_xla.core.xla_model as xm
import torch_xla.experimental.xla_sharding as xs
from torch import nn
from torch_xla.amp import autocast, GradScaler, syncfree
from torch_xla.experimental.xla_sharding import Mesh


class SimpleLinear(nn.Module):

  def __init__(self):
    super(SimpleLinear, self).__init__()
    self.fc0 = nn.Linear(1024, 1024)
    self.fc1 = nn.Linear(1024, 1024)
    self.fc2 = nn.Linear(1024, 1024)
    self.fc3 = nn.Linear(1024, 1024)
    self.relu = nn.ReLU()


  def forward(self, x):
    h = self.relu(self.fc0(x))
    y = self.relu(self.fc1(h) + x)
    z = self.fc3(y)
    return z

def run():
    device = xm.xla_device()
    model = SimpleLinear().to(device)

    num_devices = 4
    devices_ids = np.arange(num_devices)
    row_mesh = Mesh(devices_ids, (num_devices, 1), ('x', 'y'))
    col_mesh = Mesh(devices_ids, (1, num_devices), ('x', 'y'))

    xs.mark_sharding(model.fc1.weight, col_mesh, (0, 1))
    xs.mark_sharding(model.fc2.weight, row_mesh, (0, 1))

    optimizer = syncfree.AdamW(model.parameters(), lr=0.0, betas=(0.9, 0.999), eps=1e-8)
    scaler = GradScaler()

    for step in range(5):
        dummy_data = torch.zeros(64, 1024, device=device)
        x = dummy_data
        with autocast():
            x = model(x)
            dummy_loss = x.sum()
            scaler.scale(dummy_loss).backward()
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            xm.mark_step()
        print(step)


if __name__ == "__main__":
    run()

XLA_USE_SPMD=1 torchrun --nproc_per_node=4 examples/test_spmd_amp.py