def METHOD_NAME(strategy):
    rank = strategy.local_rank
    world_size = strategy.num_processes
    tensor1 = torch.ones(8, requires_grad=True)
    tensor2 = torch.ones((8, 16, 32), requires_grad=True)
    tensor1_gathered = strategy.all_gather(tensor1, sync_grads=True)
    tensor2_gathered = strategy.all_gather(tensor2, sync_grads=True)
    tensor1_gathered = tensor1_gathered * rank
    tensor2_gathered = tensor2_gathered * rank
    tensor1_gathered.sum().backward()
    tensor2_gathered.sum().backward()
    grad1 = torch.zeros_like(tensor1.grad).fill_(torch.arange(world_size).sum().float())
    grad2 = torch.zeros_like(tensor2.grad).fill_(torch.arange(world_size).sum().float())
    assert torch.allclose(grad1, tensor1.grad)
    assert torch.allclose(grad2, tensor2.grad)