code: def test_compile_selective_checkpoint_custom_rule(self):
      def _get_custom_policy(meta):
          no_recompute_list = [
              torch.ops.aten.mm.default,
          ]
  
          def _custom_policy(mode, func, *args, **kwargs):
              mm_count_key = f"{mode}_mm_count"
              if mm_count_key not in meta:
                  meta[mm_count_key] = 0
              if func == torch.ops.aten.mm.default:
                  meta[mm_count_key] += 1
              # Saves output of all compute ops, except second mm
              # (i.e. we will hint the partitioner to recompute second mm in backward pass)
              return func in no_recompute_list and not (
                  func == torch.ops.aten.mm.default
                  and meta[mm_count_key] == 2
              )
  
          return _custom_policy
  
      def selective_checkpointing_context_fn():
          meta = {}
          return _pt2_selective_checkpoint_context_fn_gen(
              _get_custom_policy(meta)
          )
  
      def gn(x, y):
          return torch.sigmoid(
              torch.sigmoid(torch.matmul(torch.matmul(x, y) * y, y) * y)
          )
  
      def fn(x, y):
          return torch.utils.checkpoint.checkpoint(
              gn,
              x,
              y,
              use_reentrant=False,
              context_fn=selective_checkpointing_context_fn,
          )
  
      x = torch.randn(4, 4, requires_grad=True, device="cuda")
      y = torch.randn(4, 4, requires_grad=True, device="cuda")
  
      fw_compiler = functools.partial(
          count_ops,
          freq=2,
          op=torch.ops.aten.mm.default,
      )
      bw_compiler = functools.partial(
          count_ops,
          # Q: How do we come to this number 4?
          # A: We have 2 matmuls in the forward pass, each matmul contributes 2 `mm` ops in the backward pass,
          # so we have at least 4 `mm` ops in backward pass. It's "at least" because whether second matmul in
          # the forward pass is recomputed in the backward pass is up to the partitioner to decide.
          freq_ge=4,
          op=torch.ops.aten.mm.default,
      )
      backend = aot_autograd(
          fw_compiler=fw_compiler,
          bw_compiler=bw_compiler,
          partition_fn=min_cut_rematerialization_partition,
      )
      self._validate(fn, backend, x, y)
      self._compare_orig_and_checkpointed_fns(gn, fn, x, y)
  
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 17 for calling the ActivationCheckpointingViaTagsTests obj's underlying member method code with example usage: obj.test_compile_selective_checkpoint_custom_rule() and returns: _custom_policy, _pt2_selective_checkpoint_context_fn_gen, torch, torch, unknown