class_docstr: This operator is supposed to be used only with torch.compile stack. This
  accepts a Fx graph module which needs to be checkpointed. This operator adds
  "recomputable" tag to the nodes of the Fx graph that should be recomputed.
  
  The goal is to:
  1. Avoid using Dynamo to trace through saved tensor hooks.
  2. For selective checkpointing case, let AOTAutograd trace through
     saved tensor hooks but has special logic with TorchDispatchMode to override
     the usual saved_tensor_hooks fn logic in order to tag the nodes.
  3. Rely on the partitioners to actually duplicate the nodes.
  This sits well in the torch.compile stack, because by the time graph
  reaches partitioner, inductor has already run its functionalization of rng
  ops (by setting fixed seed for each random op, see `replace_random_passes`).
  Therefore, the duplication of nodes, by design, respects the rng states in
  the forward and recomputed forward in backward.
class_docstr_tok: ["This","operator","is","supposed","to","be","used","only","with","torch.compile","stack",".","This","accepts","a","Fx","graph","module","which","needs","to","be","checkpointed",".","This","operator","adds","''","recomputable","''","tag","to","the","nodes","of","the","Fx","graph","that","should","be","recomputed",".","The","goal","is","to",":","1",".","Avoid","using","Dynamo","to","trace","through","saved","tensor","hooks",".","2",".","For","selective","checkpointing","case",",","let","AOTAutograd","trace","through","saved","tensor","hooks","but","has","special","logic","with","TorchDispatchMode","to","override","the","usual","saved_tensor_hooks","fn","logic","in","order","to","tag","the","nodes",".","3",".","Rely","on","the","partitioners","to","actually","duplicate","the","nodes",".","This","sits","well","in","the","torch.compile","stack",",","because","by","the","time","graph","reaches","partitioner",",","inductor","has","already","run","its","functionalization","of","rng","ops","(","by","setting","fixed","seed","for","each","random","op",",","see","`","replace_random_passes","`",")",".","Therefore",",","the","duplication","of","nodes",",","by","design",",","respects","the","rng","states","in","the","forward","and","recomputed","forward","in","backward","."]
code: def tag_nodes(self, gmod):
      unique_graph_id = next(uid)
      for node in gmod.graph.nodes:
          if node.op in ("call_function", "call_method", "call_module"):
              node.meta["recompute"] = unique_graph_id
      return gmod
  
file_path: pytorch/torch/_higher_order_ops/wrap.py
imports: import inspect
  import logging
  import torch
  from torch._ops import HigherOrderOperator
  from torch.utils.checkpoint import checkpoint, uid
  import torch._dynamo.config
label_desc: Use image node_id 3 for calling the TagActivationCheckpoint obj's underlying member method code with example usage: obj.tag_nodes(gmod) and returns: gmod