code: def check_forward_full_state_property(
      metric_class: Metric,
      init_args: Optional[Dict[str, Any]] = None,
      input_args: Optional[Dict[str, Any]] = None,
      num_update_to_compare: Sequence[int] = [10, 100, 1000],
      reps: int = 5,
  ) -> None:
      """Check if the new ``full_state_update`` property works as intended.
  
      This function checks if the property can safely be set to ``False`` which will for most metrics results in a
      speedup when using ``forward``.
  
      Args:
          metric_class: metric class object that should be checked
          init_args: dict containing arguments for initializing the metric class
          input_args: dict containing arguments to pass to ``forward``
          num_update_to_compare: if we successfully detect that the flag is safe to set to ``False``
              we will run some speedup test. This arg should be a list of integers for how many
              steps to compare over.
          reps: number of repetitions of speedup test
  
      Example (states in ``update`` are independent, save to set ``full_state_update=False``)
          >>> from torchmetrics.classification import MulticlassConfusionMatrix
          >>> check_forward_full_state_property(  # doctest: +SKIP
          ...     MulticlassConfusionMatrix,
          ...     init_args = {'num_classes': 3},
          ...     input_args = {'preds': torch.randint(3, (100,)), 'target': torch.randint(3, (100,))},
          ... )
          Full state for 10 steps took: ...
          Partial state for 10 steps took: ...
          Full state for 100 steps took: ...
          Partial state for 100 steps took: ...
          Full state for 1000 steps took: ...
          Partial state for 1000 steps took: ...
          Recommended setting `full_state_update=False`
  
      Example (states in ``update`` are dependent meaning that ``full_state_update=True``):
          >>> from torchmetrics.classification import MulticlassConfusionMatrix
          >>> class MyMetric(MulticlassConfusionMatrix):
          ...     def update(self, preds, target):
          ...         super().update(preds, target)
          ...         # by construction make future states dependent on prior states
          ...         if self.confmat.sum() > 20:
          ...             self.reset()
          >>> check_forward_full_state_property(
          ...     MyMetric,
          ...     init_args = {'num_classes': 3},
          ...     input_args = {'preds': torch.randint(3, (10,)), 'target': torch.randint(3, (10,))},
          ... )
          Recommended setting `full_state_update=True`
  
      """
      init_args = init_args or {}
      input_args = input_args or {}
  
      class FullState(metric_class):
          full_state_update = True
  
      class PartState(metric_class):
          full_state_update = False
  
      fullstate = FullState(**init_args)
      partstate = PartState(**init_args)
  
      equal = True
      try:  # if it fails, the code most likely need access to the full state
          for _ in range(num_update_to_compare[0]):
              equal = equal & _allclose_recursive(
                  fullstate(**input_args), partstate(**input_args)
              )
      except RuntimeError:
          equal = False
      res1 = fullstate.compute()
      try:  # if it fails, the code most likely need access to the full state
          res2 = partstate.compute()
      except RuntimeError:
          equal = False
      equal = equal & _allclose_recursive(res1, res2)
  
      if (
          not equal
      ):  # we can stop early because the results did not match
          print("Recommended setting `full_state_update=True`")
          return
  
      # Do timings
      res = torch.zeros(2, len(num_update_to_compare), reps)
      for i, metric in enumerate([fullstate, partstate]):
          for j, t in enumerate(num_update_to_compare):
              for r in range(reps):
                  start = perf_counter()
                  for _ in range(t):
                      _ = metric(**input_args)
                  end = perf_counter()
                  res[i, j, r] = end - start
                  metric.reset()
  
      mean = torch.mean(res, -1)
      std = torch.std(res, -1)
  
      for t in range(len(num_update_to_compare)):
          print(
              f"Full state for {num_update_to_compare[t]} steps took: {mean[0, t]}+-{std[0, t]:0.3f}"
          )
          print(
              f"Partial state for {num_update_to_compare[t]} steps took: {mean[1, t]:0.3f}+-{std[1, t]:0.3f}"
          )
  
      faster = (
          mean[1, -1] < mean[0, -1]
      ).item()  # if faster on average, we recommend upgrading
      print(f"Recommended setting `full_state_update={not faster}`")
      return
  
imports: import logging
  import multiprocessing
  import os
  from functools import partial
  from time import perf_counter
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, no_type_check
  from unittest.mock import Mock
  import torch
  from torch import Tensor
  from torchmetrics.metric import Metric
  from torchmetrics.utilities.data import select_topk, to_onehot
  from torchmetrics.utilities.enums import DataType
label_desc: Use image node_id 17 for calling a global function with example usage: check_forward_full_state_property(metric_class, init_args, input_args, num_update_to_compare, reps) without return types