code: def _functional_test(
      preds: Union[Tensor, list],
      target: Union[Tensor, list, List[Dict[str, Tensor]]],
      metric_functional: Callable,
      reference_metric: Callable,
      metric_args: Optional[dict] = None,
      atol: float = 1e-8,
      device: str = "cpu",
      fragment_kwargs: bool = False,
      **kwargs_update: Any,
  ):
      """Comparison between functional metric and reference metric.
  
      Args:
          preds: torch tensor with predictions
          target: torch tensor with targets
          metric_functional: metric functional that should be tested
          reference_metric: callable function that is used for comparison
          metric_args: dict with additional arguments used for class initialization
          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
          kwargs_update: Additional keyword arguments that will be passed with preds and
              target when running update on the metric.
  
      """
      p_size = (
          preds.shape[0] if isinstance(preds, Tensor) else len(preds)
      )
      t_size = (
          target.shape[0] if isinstance(target, Tensor) else len(target)
      )
      assert p_size == t_size, f"different sizes {p_size} and {t_size}"
      num_batches = p_size
      metric_args = metric_args or {}
      metric = partial(metric_functional, **metric_args)
  
      # move to device
      if isinstance(preds, Tensor):
          preds = preds.to(device)
      if isinstance(target, Tensor):
          target = target.to(device)
      elif isinstance(target, list):
          for i, target_dict in enumerate(target):
              if isinstance(target_dict, dict):
                  for k in target_dict:
                      if isinstance(target_dict[k], Tensor):
                          target[i][k] = target_dict[k].to(device)
  
      kwargs_update = {
          k: v.to(device) if isinstance(v, Tensor) else v
          for k, v in kwargs_update.items()
      }
  
      for i in range(num_batches // 2):
          extra_kwargs = {
              k: v[i] if isinstance(v, Tensor) else v
              for k, v in kwargs_update.items()
          }
          tm_result = metric(preds[i], target[i], **extra_kwargs)
          extra_kwargs = {
              k: v.cpu() if isinstance(v, Tensor) else v
              for k, v in (
                  extra_kwargs if fragment_kwargs else kwargs_update
              ).items()
          }
          ref_result = reference_metric(
              preds[i].cpu() if isinstance(preds, Tensor) else preds[i],
              target[i].cpu()
              if isinstance(target, Tensor)
              else target[i],
              **extra_kwargs,
          )
          # assert it is the same
          _assert_allclose(tm_result, ref_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 5 for calling a global function with example usage: _functional_test(preds, target, metric_functional, reference_metric, metric_args, atol, device, fragment_kwargs) without return types