code: def _wrap_reduction(
      p: Tensor,
      q: Tensor,
      log_prob: bool,
      reduction: Optional[str] = "mean",
  ):
      if log_prob:
          p = p.softmax(dim=-1)
          q = q.softmax(dim=-1)
      res = entropy(p, q, axis=1)
      if reduction == "mean":
          return np.mean(res)
      if reduction == "sum":
          return np.sum(res)
      return res
  
imports: from functools import partial
  from typing import NamedTuple, Optional
  import numpy
  import pytest
  import torch
  from scipy.stats import entropy
  from torch import Tensor
  from torchmetrics.functional.regression.kl_divergence import kl_divergence
  from torchmetrics.regression.kl_divergence import KLDivergence
  from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from unittests import BATCH_SIZE, EXTRA_DIM, NUM_BATCHES
  from unittests.helpers import seed_all
  from unittests.helpers.testers import MetricTester
label_desc: Use image node_id 1 for calling a global function with example usage: _wrap_reduction(p, q, log_prob, reduction) and returns: res, np, np