code: def _compare_orig_and_checkpointed_fns(
      self, orig_fn, checkpointed_fn, *args, fullgraph=True
  ):
      # The original version and the checkpointed version of the same function
      # should produce the same outputs and the same gradients under torch.compile.
  
      # Run original version
      cloned_args_orig_fn = []
      for arg in args:
          cloned_args_orig_fn.append(
              arg.clone().detach().requires_grad_(arg.requires_grad)
          )
      torch.manual_seed(0)
      compiled_orig_fn = torch.compile(
          orig_fn, fullgraph=fullgraph, backend="inductor"
      )
      result_orig_fn = compiled_orig_fn(*cloned_args_orig_fn)
      result_orig_fn.sum().backward()
  
      # Run checkpointed version
      cloned_args_checkpointed_fn = []
      for arg in args:
          cloned_args_checkpointed_fn.append(
              arg.clone().detach().requires_grad_(arg.requires_grad)
          )
      torch.manual_seed(0)
      compiled_checkpointed_fn = torch.compile(
          checkpointed_fn, fullgraph=fullgraph, backend="inductor"
      )
      result_checkpointed_fn = compiled_checkpointed_fn(
          *cloned_args_checkpointed_fn
      )
      result_checkpointed_fn.sum().backward()
  
      # Check that outputs and gradients are equal
      self.assertEqual(
          result_orig_fn,
          result_checkpointed_fn,
          msg="Output mismatch between the original version and the checkpointed version of the same function",
      )
      for cloned_arg_orig_fn, cloned_arg_checkpointed_fn in zip(
          cloned_args_orig_fn, cloned_args_checkpointed_fn
      ):
          self.assertEqual(
              cloned_arg_orig_fn.grad,
              cloned_arg_checkpointed_fn.grad,
              msg="Gradient mismatch between the original version and the checkpointed version of the same function",
          )
  
imports: import functools
  import unittest
  from importlib import import_module
  import torch
  import torch._dynamo.config
  import torch._dynamo.test_case
  import torch._functorch.config
  import torch.utils.checkpoint
  from functorch.compile import min_cut_rematerialization_partition
  from torch._dynamo.backends.common import aot_autograd
  from torch._dynamo.testing import CompileCounterWithBackend
  from torch._higher_order_ops.wrap import tag_activation_checkpoint
  from torch.testing._internal.common_utils import IS_WINDOWS, skipIfRocm
  from torch.testing._internal.inductor_utils import HAS_CUDA
  from torch.testing._internal.two_tensor import TwoTensor
  from torch.utils.checkpoint import _pt2_selective_checkpoint_context_fn_gen, checkpoint
label_desc: Use image node_id 2 for calling the ActivationCheckpointingViaTagsTests obj's underlying member method code with example usage: obj._compare_orig_and_checkpointed_fns(orig_fn, checkpointed_fn) without return types