code: def _class_test(
      rank: int,
      world_size: int,
      preds: Union[Tensor, list, List[Dict[str, Tensor]]],
      target: Union[Tensor, list, List[Dict[str, Tensor]]],
      metric_class: Metric,
      reference_metric: Callable,
      dist_sync_on_step: bool,
      metric_args: Optional[dict] = None,
      check_dist_sync_on_step: bool = True,
      check_batch: bool = True,
      atol: float = 1e-8,
      device: str = "cpu",
      fragment_kwargs: bool = False,
      check_scriptable: bool = True,
      check_state_dict: bool = True,
      **kwargs_update: Any,
  ):
      """Comparison between class metric and reference metric.
  
      Args:
          rank: rank of current process
          world_size: number of processes
          preds: torch tensor with predictions
          target: torch tensor with targets
          metric_class: metric class that should be tested
          reference_metric: callable function that is used for comparison
          dist_sync_on_step: bool, if true will synchronize metric state across
              processes at each ``forward()``
          metric_args: dict with additional arguments used for class initialization
          check_dist_sync_on_step: bool, if true will check if the metric is also correctly
              calculated per batch and per device (and not just at the end)
          check_batch: bool, if true will check if the metric is also correctly
              calculated across devices for each batch (and not just at the end)
          atol: absolute tolerance used for comparison of results
          device: determine which device to run on, either 'cuda' or 'cpu'
          fragment_kwargs: whether tensors in kwargs should be divided as `preds` and `target` among processes
          check_scriptable: bool indicating if metric should also be tested if it can be scripted
          check_state_dict: bool indicating if metric should be tested that its state_dict by default is empty
          kwargs_update: Additional keyword arguments that will be passed with preds and
              target when running update on the metric.
  
      """
      assert len(preds) == len(target)
      num_batches = len(preds)
      assert (
          num_batches % world_size == 0
      ), "Number of batches must be divisible by world_size"
  
      if not metric_args:
          metric_args = {}
  
      # Instantiate metric
      metric = metric_class(
          dist_sync_on_step=dist_sync_on_step, **metric_args
      )
      with pytest.raises(RuntimeError):
          metric.is_differentiable = not metric.is_differentiable
      with pytest.raises(RuntimeError):
          metric.higher_is_better = not metric.higher_is_better
  
      # check that the metric is scriptable
      if check_scriptable:
          torch.jit.script(metric)
  
      # check that metric can be cloned
      clone = metric.clone()
      assert (
          clone is not metric
      ), "Clone is not a different object than the metric"
      assert type(clone) == type(
          metric
      ), "Type of clone did not match metric type"
  
      # move to device
      metric = metric.to(device)
      preds = apply_to_collection(preds, Tensor, lambda x: x.to(device))
      target = apply_to_collection(
          target, Tensor, lambda x: x.to(device)
      )
  
      kwargs_update = {
          k: v.to(device) if isinstance(v, Tensor) else v
          for k, v in kwargs_update.items()
      }
  
      # verify metrics work after being loaded from pickled state
      pickled_metric = pickle.dumps(metric)
      metric = pickle.loads(pickled_metric)
  
      for i in range(rank, num_batches, world_size):
          batch_kwargs_update = {
              k: v[i] if isinstance(v, Tensor) else v
              for k, v in kwargs_update.items()
          }
  
          # compute batch stats and aggregate for global stats
          batch_result = metric(
              preds[i], target[i], **batch_kwargs_update
          )
  
          if (
              metric.dist_sync_on_step
              and check_dist_sync_on_step
              and rank == 0
          ):
              if isinstance(preds, Tensor):
                  ddp_preds = torch.cat(
                      [preds[i + r] for r in range(world_size)]
                  ).cpu()
              else:
                  ddp_preds = _flatten(
                      [preds[i + r] for r in range(world_size)]
                  )
              if isinstance(target, Tensor):
                  ddp_target = torch.cat(
                      [target[i + r] for r in range(world_size)]
                  ).cpu()
              else:
                  ddp_target = _flatten(
                      [target[i + r] for r in range(world_size)]
                  )
              ddp_kwargs_upd = {
                  k: torch.cat(
                      [v[i + r] for r in range(world_size)]
                  ).cpu()
                  if isinstance(v, Tensor)
                  else v
                  for k, v in (
                      kwargs_update
                      if fragment_kwargs
                      else batch_kwargs_update
                  ).items()
              }
              ref_batch_result = reference_metric(
                  ddp_preds, ddp_target, **ddp_kwargs_upd
              )
              if isinstance(batch_result, dict):
                  for key in batch_result:
                      _assert_allclose(
                          batch_result,
                          ref_batch_result[key].numpy(),
                          atol=atol,
                          key=key,
                      )
              else:
                  _assert_allclose(
                      batch_result, ref_batch_result, atol=atol
                  )
  
          elif check_batch and not metric.dist_sync_on_step:
              batch_kwargs_update = {
                  k: v.cpu() if isinstance(v, Tensor) else v
                  for k, v in (
                      batch_kwargs_update
                      if fragment_kwargs
                      else kwargs_update
                  ).items()
              }
              preds_ = (
                  preds[i].cpu()
                  if isinstance(preds, Tensor)
                  else preds[i]
              )
              target_ = (
                  target[i].cpu()
                  if isinstance(target, Tensor)
                  else target[i]
              )
              ref_batch_result = reference_metric(
                  preds_, target_, **batch_kwargs_update
              )
              if isinstance(batch_result, dict):
                  for key in batch_result:
                      _assert_allclose(
                          batch_result,
                          ref_batch_result[key].numpy(),
                          atol=atol,
                          key=key,
                      )
              else:
                  _assert_allclose(
                      batch_result, ref_batch_result, atol=atol
                  )
  
      # check that metrics are hashable
      assert hash(metric), repr(metric)
  
      # assert that state dict is empty
      if check_state_dict:
          assert metric.state_dict() == {}
  
      # check on all batches on all ranks
      result = metric.compute()
      if isinstance(result, dict):
          for key in result:
              _assert_tensor(result, key=key)
      else:
          _assert_tensor(result)
  
      if isinstance(preds, Tensor):
          total_preds = torch.cat(
              [preds[i] for i in range(num_batches)]
          ).cpu()
      else:
          total_preds = [item for sublist in preds for item in sublist]
      if isinstance(target, Tensor):
          total_target = torch.cat(
              [target[i] for i in range(num_batches)]
          ).cpu()
      elif (
          isinstance(target, list)
          and len(target) > 0
          and isinstance(target[0], dict)
      ):
          total_target = {
              k: torch.cat([t[k] for t in target]) for k in target[0]
          }
      else:
          total_target = [
              item for sublist in target for item in sublist
          ]
  
      total_kwargs_update = {
          k: torch.cat([v[i] for i in range(num_batches)]).cpu()
          if isinstance(v, Tensor)
          else v
          for k, v in kwargs_update.items()
      }
      sk_result = reference_metric(
          total_preds, total_target, **total_kwargs_update
      )
  
      # assert after aggregation
      if isinstance(sk_result, dict):
          for key in sk_result:
              _assert_allclose(
                  result, sk_result[key].numpy(), atol=atol, key=key
              )
      else:
          _assert_allclose(result, sk_result, atol=atol)
  
imports: import pickle
  import sys
  from copy import deepcopy
  from functools import partial
  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
  import numpy
  import pytest
  import torch
  from lightning_utilities import apply_to_collection
  from torch import Tensor, tensor
  from torchmetrics import Metric
  from torchmetrics.utilities.data import _flatten
  from unittests import NUM_PROCESSES
label_desc: Use image node_id 4 for calling a global function with example usage: _class_test(rank, world_size, preds, target, metric_class, reference_metric, dist_sync_on_step, metric_args, check_dist_sync_on_step, check_batch, atol, device, fragment_kwargs, check_scriptable, check_state_dict) without return types