code: def wrap_optimizer_init(fn: Callable):
      """Wraps the __init__ function of OptimizerV2 objects and stores the info in the default GraphItem."""
  
      def wrapper(*args, **kwargs):
          # args[0] should be `self`, which is an object of type == optimizer class
          containing_class = type(args[0])
          class_name = containing_class.__name__
  
          # For calls like super(AdamWeightDecay, self).__init__(*args, **kwargs), the containing_class.__name__
          # returns the current class (AdamWeightDecay) instead of the parent class (Adam).
          # Avoid patching this pattern by checking fn.__qualname__.
          if not fn.__qualname__.startswith(class_name):
              return fn(*args, **kwargs)
          if _default_graph_item and kwargs.pop("update", True):
              _default_graph_item.extend_optimizer_info(
                  containing_class, *args, **kwargs
              )
              logging.debug(
                  "Registered optimizer: {} \nwith args: {} \nkwargs: {}".format(
                      class_name, args, kwargs
                  )
              )
          return fn(*args, **kwargs)
  
      return wrapper
  
imports: import contextlib
  import copy
  import functools
  from collections import defaultdict
  from typing import Union, Callable
  from google.protobuf.any_pb2 import Any
  from tensorflow.core.framework.graph_pb2 import GraphDef
  from tensorflow.core.framework.variable_pb2 import VariableDef
  from tensorflow.core.protobuf.saver_pb2 import SaverDef
  from tensorflow.python.framework import ops
  from tensorflow.python.framework.importer import import_graph_def
  from tensorflow.python.ops.resource_variable_ops import _from_proto_fn
  from tensorflow.python.ops.variables import Variable
  from autodist.const import COLOCATION_PREFIX
  from autodist.kernel.common import op_info
  from autodist.kernel.common.utils import parse_name_scope, get_op_name
  from autodist.proto import graphitem_pb2
  from autodist.utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: wrap_optimizer_init(fn) and returns: wrapper, fn, fn