code: def tpu_all_gather_fn(strategy):
      with pytest.raises(
          NotImplementedError, match="only implemented for tensors"
      ):
          strategy.all_gather([1])
  
      for sync_grads in (True, False):
          tensor = torch.tensor(1.0, requires_grad=True)
          result = strategy.all_gather(tensor, sync_grads=sync_grads)
          summed = result.sum()
          assert (
              summed.device.type == "cpu"
          )  # the original device is preserved
          assert torch.equal(
              summed,
              torch.tensor(strategy.world_size, dtype=torch.float32),
          )
          if not _XLA_GREATER_EQUAL_2_1:
              summed.backward()
          if sync_grads:
              if _XLA_GREATER_EQUAL_2_1:
                  # in 2.1, sync_grads=False makes it so that you cannot call .backward even if it originally had set
                  # requires_grad=True
                  summed.backward()
              assert torch.equal(tensor.grad, torch.tensor(1.0))
          else:
              # As gradients are not synced, the original tensor will not have gradients.
              assert tensor.grad is None
  
imports: import os
  from functools import partial
  from unittest import mock
  from unittest.mock import ANY, MagicMock, Mock
  import pytest
  import torch
  from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1, XLAAccelerator
  from lightning.fabric.strategies import XLAStrategy
  from lightning.fabric.strategies.launchers.xla import _XLALauncher
  from lightning.fabric.utilities.distributed import ReduceOp
  from lightning.fabric.utilities.seed import seed_everything
  from torch.utils.data import DataLoader
  from tests_fabric.helpers.models import RandomDataset
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 8 for calling a global function with example usage: tpu_all_gather_fn(strategy) without return types