class_docstr: Implementation for working with TF2.x Graph Mode.
docstr: Experimental interface similar to :tf_main:`tf.function <function>`.
docstr_tok: ["Experimental","interface","similar","to",":","tf_main",":","`","tf.function","<","function",">","`","."]
code: def function(self, fn):
      """Experimental interface similar to :tf_main:`tf.function <function>`."""
      _cache = self._cache
      _build_fn = self._build_fn
  
      def wrapper(*args, **kwargs):
          # we first assume one fn only build one type of graph
          cache_id = hash(_V2Graph._CacheKey(fn))
          cached = cache_id in _cache
  
          # At the first run of the training function
          if not cached:
              if _cache:
                  raise NotImplementedError(
                      "AutoDist currently only stably supports "
                      "one 'autodist.function' across the scope."
                  )
              # Cache the runner
              _cache[cache_id] = _build_fn(fn, *args, **kwargs)
              atexit.register(lambda: _cache.pop(cache_id))
          return _cache[cache_id](*args, **kwargs)
  
      return wrapper
  
imports: import atexit
  import os
  from collections import namedtuple
  import numpy
  from tensorflow.python.eager import context
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.util import tf_contextlib
  from autodist.cluster import Cluster, SSHCluster
  from autodist.const import ENV
  from autodist.coordinator import Coordinator
  from autodist.graph_item import GraphItem
  from autodist.kernel.device.resolver import DeviceResolver
  from autodist.kernel.graph_transformer import GraphTransformer
  from autodist.patch import PatchTensorFlow
  from autodist.remapper import Remapper
  from autodist.resource_spec import ResourceSpec
  from autodist.runner import WrappedSession
  from autodist.strategy import base
  from autodist.strategy.ps_lb_strategy import PSLoadBalancing
  from autodist.utils import logging
label_desc: Use image node_id 5 for calling the _V2Graph obj's underlying member method code with example usage: obj.function(fn) and returns: wrapper, _cache