import torch
import torch.utils.checkpoint
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--xla", type=int, required=True)
args = parser.parse_args()

if args.xla:
    import torch_xla.core.xla_model as xm
    import torch_xla.distributed.xla_multiprocessing as xmp
    from torch_xla.distributed.fsdp import checkpoint_module

to_save = []
def save_output(output):
    to_save.append(output.detach().cpu())

class Model(torch.nn.Module):
    def __init__(self, args):
        super().__init__()
        self.x = torch.nn.Linear(128,128)
        self.dropout = torch.nn.Dropout(p=0.1)
        self.args = args
    
    def forward(self, inp):
        x = self.x(inp)
        output = self.dropout(x)
        if self.args.xla:
            xm.add_step_closure(save_output, args=(output,), run_async=False)
        else:
            save_output(output)
        return output

def main(args):
    if args.xla:
        device = xm.xla_device()
    else:
        device = 0
        torch.cuda.set_device(device)

    model = Model(args)
    model.to(device)

    if args.xla:
        model = checkpoint_module(model)

    _input = torch.randn(128, 128, requires_grad=True)
    _input = _input.to(device)

    if not args.xla:
        output = torch.utils.checkpoint.checkpoint(model, _input)
    else:
        output = model(_input)
    output = torch.sum(output)
    output.backward()
    if args.xla:
        xm.mark_step()
    same_output = torch.allclose(to_save[0], to_save[1])
    print(f"xla {args.xla} same_output {same_output}")

if __name__ == "__main__":    
    main(args)

python test_dropout_simple.py --xla 0
xla 0 same_output True

python test_dropout_simple.py --xla 1
xla 0 same_output False