code: def tpu_reduce_fn(strategy):
      with pytest.raises(ValueError, match="XLAStrategy only supports"):
          strategy.all_reduce(1, reduce_op="undefined")
  
      with pytest.raises(ValueError, match="XLAStrategy only supports"):
          strategy.all_reduce(1, reduce_op=ReduceOp.MAX)
  
          # it is faster to loop over here than to parameterize the test
          for reduce_op in ("mean", "AVG", "sum", ReduceOp.SUM):
              result = strategy.all_reduce(1, reduce_op=reduce_op)
              if isinstance(reduce_op, str) and reduce_op.lower() in (
                  "mean",
                  "avg",
              ):
                  assert result.item() == 1
              else:
                  assert result.item() == 8
  
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 5 for calling a global function with example usage: tpu_reduce_fn(strategy) without return types