code: def tpu_sync_module_states_fn(sync_module_states, strategy):
      seed_everything(
          strategy.local_rank
      )  # force the model to have different weights across ranks
      model = torch.nn.Linear(1, 1).to(strategy.root_device)
      model = strategy.setup_module(model)
      gathered = strategy.all_gather(model.weight)
      for t in gathered[1:]:
          if sync_module_states:
              assert torch.equal(gathered[0], t)
          else:
              assert not torch.equal(gathered[0], t)
  
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 10 for calling a global function with example usage: tpu_sync_module_states_fn(sync_module_states, strategy) without return types