code: def broadcast_on_tpu_fn(strategy):
      # test broadcasting a tensor
      obj = torch.tensor(strategy.global_rank)
      assert obj.device.type == "cpu"
      # In PjRT, the local rank and global rank have no solid relation.
      # global rank may not even be contiguous on a host, because it depends on the 3D mesh structure that is formed by
      # the TPUs on all hosts in a pod. So checking a different src is not reliable
      # https://github.com/pytorch/xla/blob/v2.0.0/torch_xla/experimental/pjrt.py#L161-L163
      src = 0
      result = strategy.broadcast(obj, src)
      assert result.item() == src
      assert (
          result.device.type == "cpu"
      )  # the original device is preserved
  
      # test broadcasting an arbitrary object
      tensor = torch.tensor(
          strategy.global_rank,
          device=strategy.root_device,
          dtype=torch.bfloat16,
      )
      obj = ("ver_0.5", "logger_name", strategy.global_rank, tensor)
      result = strategy.broadcast(obj, src=src)
      assert result == ("ver_0.5", "logger_name", src, ANY)
      assert (
          result[3].device.type == "xla"
      )  # the original device is preserved
      assert result[3].dtype == torch.bfloat16
  
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 3 for calling a global function with example usage: broadcast_on_tpu_fn(strategy) without return types