code: def wrap_optimizer_apply_gradient(fn: Callable):
      """Wraps the apply_gradients function of OptimizerV2 objects and stores the info in the default GraphItem."""
  
      # Signature for apply_gradients
      # apply_gradients(self, grads_and_vars, name=None)
      def wrapper(*args, **kwargs):
          # Assume grads_and_vars is an iterable of tuples
          # Materialize here because in case it's a generator, we need to be able to iterate multiple times
          grads_and_vars = list(kwargs.get("grads_and_vars") or args[1])
          grads, variables = map(list, zip(*grads_and_vars))
          if _default_graph_item and kwargs.pop("update", True):
              _default_graph_item.extend_gradient_info(grads, variables)
              logging.debug(
                  "Registered grads: \n {} with targets: \n {}".format(
                      grads, variables
                  )
              )
          args = (
              args[0],
              grads_and_vars,
          )  # Replace possible generator with definite list
          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 4 for calling a global function with example usage: wrap_optimizer_apply_gradient(fn) and returns: wrapper, fn