code: def test_xla_fsdp_train_save_load(
      tmp_path, use_auto_wrap_policy, state_dict_type, sequential_save
  ):
      """Test XLAFSDP training, saving and loading checkpoint (both full and sharded)."""
      from torch_xla.distributed.fsdp.wrap import always_wrap_policy
  
      policy = always_wrap_policy if use_auto_wrap_policy else None
      strategy = XLAFSDPStrategy(
          auto_wrap_policy=policy,
          state_dict_type=state_dict_type,
          sequential_save=sequential_save,
      )
      fabric = Fabric(accelerator="tpu", strategy=strategy)
      fabric.launch(xla_fsdp_train_save_load, tmp_path, state_dict_type)
  
file_path: pytorch-lightning/tests/tests_fabric/strategies/test_xla_fsdp_integration.py
imports: import os
  import re
  from pathlib import Path
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.strategies import XLAFSDPStrategy
  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 4 for calling a global function with example usage: test_xla_fsdp_train_save_load(tmp_path, use_auto_wrap_policy, state_dict_type, sequential_save) without return types