code: def _perplexity_class_test(
      rank: int,
      worldsize: int,
      probs: Optional[torch.Tensor],
      logits: Optional[torch.Tensor],
      dist_sync_on_step: bool,
      metric_args: dict = {},
      check_dist_sync_on_step: bool = True,
      check_batch: bool = True,
      atol: float = 1e-8,
  ):
      """Utility function doing the actual comparison between lightning class metric
      and reference metric.
      Args:
          rank: rank of current process
          worldsize: number of processes
          probs: torch tensor with probabilities
          logits: torch tensor with logits. The function checks ``probs`` and ``logits are mutually exclusive for
              ``Perplexity`` metric.
          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 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)
      """
      # Instanciate lightning metric
      perplexity = Perplexity(
          dist_sync_on_step=dist_sync_on_step, **metric_args
      )
      if (probs is None) == (logits is None):
          with pytest.raises(ValueError):
              perplexity(probs, logits)
          return
  
      # verify perplexity works after being loaded from pickled state
      pickled_metric = pickle.dumps(perplexity)
      perplexity = pickle.loads(pickled_metric)
  
      for i in range(rank, NUM_BATCHES, worldsize):
          batch_result = perplexity(
              None if probs is None else probs[i],
              None if logits is None else logits[i],
          )
  
          if perplexity.dist_sync_on_step:
              if rank == 0:
                  if probs is not None:
                      ddp_probs = torch.stack(
                          [probs[i + r] for r in range(worldsize)]
                      )
                  else:
                      ddp_logits = torch.stack(
                          [logits[i + r] for r in range(worldsize)]
                      )
                      ddp_probs = logits_to_probs(
                          ddp_logits, is_binary=False
                      )
                  sk_batch_result = reference_perplexity_func(ddp_probs)
                  # assert for dist_sync_on_step
                  if check_dist_sync_on_step:
                      assert np.allclose(
                          batch_result.numpy(),
                          sk_batch_result,
                          atol=atol,
                      )
          else:
              if probs is None:
                  p = logits_to_probs(logits[i], is_binary=False)
              else:
                  p = probs[i]
              sk_batch_result = reference_perplexity_func(p)
              # assert for batch
              if check_batch:
                  assert np.allclose(
                      batch_result.numpy(), sk_batch_result, atol=atol
                  )
  
      assert (probs is None) != (logits is None)
      # check on all batches on all ranks
      result = perplexity.compute()
      assert isinstance(result, torch.Tensor)
  
      if probs is None:
          probs = logits_to_probs(logits, is_binary=False)
      sk_result = reference_perplexity_func(probs)
  
      # assert after aggregation
      assert np.allclose(result.numpy(), sk_result, atol=atol)
  
imports: import os
  import pickle
  import sys
  from functools import partial
  from typing import Callable, Optional
  import numpy
  import pytest
  import torch
  from scipy.stats import entropy
  from torch.distributions.utils import logits_to_probs
  from torch.multiprocessing import Pool, set_start_method
  from torchmetrics import Metric
  from nemo.collections.common.metrics import GlobalAverageLossMetric, Perplexity
label_desc: Use image node_id 5 for calling a global function with example usage: _perplexity_class_test(rank, worldsize, probs, logits, dist_sync_on_step, metric_args, check_dist_sync_on_step, check_batch, atol) without return types