code: def test_all_gather_collection(tmpdir):
      class TestModel(BoringModel):
          on_train_epoch_end_called = False
  
          def on_train_epoch_end(self):
              losses = torch.rand(2, 2).t()
              gathered_loss = self.all_gather(
                  {
                      "losses_tensor_int": losses.int(),
                      "losses_tensor_float": losses,
                      "losses_tensor_list": [losses, losses],
                      "losses_np_ndarray": np.array([1, 2, 3]),
                      "losses_bool": [True, False],
                      "losses_float": [0.0, 1.0, 2.0],
                      "losses_int": [0, 1, 2],
                  }
              )
              assert (
                  gathered_loss["losses_tensor_int"][0].dtype
                  == torch.int32
              )
              assert (
                  gathered_loss["losses_tensor_float"][0].dtype
                  == torch.float
              )
              assert (
                  gathered_loss["losses_np_ndarray"][0].dtype
                  == torch.int64
              )
              # torch.bool can't be all_gathered
              assert (
                  gathered_loss["losses_bool"][0].dtype == torch.uint8
              )
              assert (
                  gathered_loss["losses_float"][0].dtype == torch.float
              )
              assert gathered_loss["losses_int"][0].dtype == torch.int
  
              losses_numel = losses.numel()
              assert (
                  gathered_loss["losses_tensor_int"].numel()
                  == 2 * losses_numel
              )
              assert (
                  gathered_loss["losses_tensor_float"].numel()
                  == 2 * losses_numel
              )
              assert torch.stack(
                  gathered_loss["losses_tensor_list"]
              ).shape == (2, 2, 2, 2)
              assert gathered_loss["losses_np_ndarray"].numel() == 2 * 3
              assert torch.stack(
                  gathered_loss["losses_bool"]
              ).shape == (2, 2)
              assert torch.stack(
                  gathered_loss["losses_float"]
              ).shape == (3, 2)
              assert torch.stack(gathered_loss["losses_int"]).shape == (
                  3,
                  2,
              )
              self.on_train_epoch_end_called = True
  
      model = TestModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=8,
          limit_val_batches=0,
          max_epochs=1,
          log_every_n_steps=1,
          accumulate_grad_batches=2,
          accelerator="gpu",
          devices=2,
          strategy="ddp",
          enable_progress_bar=False,
          enable_model_summary=False,
          enable_checkpointing=False,
      )
      trainer.fit(model)
      assert model.on_train_epoch_end_called
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_all_gather_grad.py
imports: import numpy
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from tests_pytorch.core.test_results import spawn_launch
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 3 for calling a global function with example usage: test_all_gather_collection(tmpdir) without return types