class_docstr: Remaps feeds and fetches for Sessions (and eager mode).
  
  Sets itself as the default remapper when in its scope
  and remaps graph elements to account for AutoDist's graph
  transformation.
  
  Feed Behavior:
  
  * We want to remap feeds to the corresponding element in the
    AutoDist-transformed graph. If the element has a polymorphic
    size (e.g. it has a `None` size for the batch dimension),
    we split the feed up across all replicas of the element.
    Otherwise, we duplicate the feed across all replicas of
    the element.
  
  Fetch Behavior:
  
  * We want to remap the potential new shards and replicas of
    a graph element into the original fetch.
  * Cases:
  
      * If original fetch exists (which is not affected by graph transformation), fetch the original.
      * Otherwise, for fetches that are train_ops, fetch them from all replicas;
      * for other fetches, only fetch it from the master replica.
  
          * For example, for partitioned vars, it corresponds to the concat one as_tensor on the first replica.
class_docstr_tok: ["Remaps","feeds","and","fetches","for","Sessions","(","and","eager","mode",")",".","Sets","itself","as","the","default","remapper","when","in","its","scope","and","remaps","graph","elements","to","account","for","AutoDist","'s","graph","transformation",".","Feed","Behavior",":","*","We","want","to","remap","feeds","to","the","corresponding","element","in","the","AutoDist-transformed","graph",".","If","the","element","has","a","polymorphic","size","(","e.g",".","it","has","a","`","None","`","size","for","the","batch","dimension",")",",","we","split","the","feed","up","across","all","replicas","of","the","element",".","Otherwise",",","we","duplicate","the","feed","across","all","replicas","of","the","element",".","Fetch","Behavior",":","*","We","want","to","remap","the","potential","new","shards","and","replicas","of","a","graph","element","into","the","original","fetch",".","*","Cases",":","*","If","original","fetch","exists","(","which","is","not","affected","by","graph","transformation",")",",","fetch","the","original",".","*","Otherwise",",","for","fetches","that","are","train_ops",",","fetch","them","from","all","replicas",";","*","for","other","fetches",",","only","fetch","it","from","the","master","replica",".","*","For","example",",","for","partitioned","vars",",","it","corresponds","to","the","concat","one","as_tensor","on","the","first","replica","."]
docstr: Switch the current global default mapper to be the current one and register conversion functions.
docstr_tok: ["Switch","the","current","global","default","mapper","to","be","the","current","one","and","register","conversion","functions","."]
code: def _set_default(self):
      """Switch the current global default mapper to be the current one and register conversion functions."""
      assert Remapper._default_remapper is None
      _autodist_fetch_fn = self._remap_fetch
      _autodist_feed_fn = self._remap_feed
      for i, expansion in enumerate(_REGISTERED_EXPANSIONS):
          (
              tensor_type,
              fetch_fn,
              feed_fn,
              feed_fn_for_partial_run,
          ) = expansion
  
          # Register nested conversion functions while keeping the function types
          # tensor_type: type
          # fetch_fn: fetch -> (List[fetch], List[fetched_val] -> final_fetched_val)
          # feed_fn: (feed, feed_val) -> List[(feed, feed_val)]
          # feed_fn_for_partial_run: feed -> List[feed]
  
          def nested_fetch_fn(fetch, fetch_fn=fetch_fn):
              """
              Two-level nested fetch function.
  
              TensorFlow ElementMapper Fetch Expansion -> AutoDist Fetch Expansion ->
              AutoDist Fetch Values Contraction -> TensorFlow ElementMapper Fetch Values Contraction
  
              Note that `fetch_fn(local_var)=fetch_fn(instant_outer_var)` is to avoid cell-var-from-loop issue.
              """
              fetches, contract_fn = fetch_fn(fetch)
              final_fetches = []
              inner_contract_fns = {}
              for f in fetches:
                  ff, fn = _autodist_fetch_fn(f)
                  i, j = len(final_fetches), len(final_fetches) + len(
                      ff
                  )
                  final_fetches.extend(ff)
                  inner_contract_fns[(i, j)] = fn
              return final_fetches, lambda fetched_vals: contract_fn(
                  [
                      fn(fetched_vals[i:j])
                      for (i, j), fn in sorted(
                          inner_contract_fns.items()
                      )
                  ]
              )
  
          def nested_feed_fn(feed, feed_val, feed_fn=feed_fn):
              feeds = feed_fn(feed, feed_val)
              final_feeds = []
              for fd, fv in feeds:
                  ff, _ = _autodist_feed_fn(fd, fv)
                  final_feeds.extend(ff)
              return final_feeds
  
          def nested_feed_fn_for_partial_run(
              feed, feed_fn_for_partial_run=feed_fn_for_partial_run
          ):
              feeds = feed_fn_for_partial_run(feed)
              final_feeds = []
              for fd in feeds:
                  ff, _ = _autodist_feed_fn(fd)
                  final_feeds.extend(ff)
              return final_feeds
  
          # Backup the _REGISTERED_EXPANSTION[i]
          Remapper._default_registered_expansions.append(expansion)
          # Write the new _REGISTERED_EXPANSTION[i]
          _REGISTERED_EXPANSIONS[i] = (
              tensor_type,
              nested_fetch_fn,
              nested_feed_fn,
              nested_feed_fn_for_partial_run,
          )
  
      Remapper._default_remapper = self
  
file_path: autodist/autodist/remapper.py
imports: import contextlib
  import numpy
  from tensorflow.python.client.session import _REGISTERED_EXPANSIONS
  from tensorflow.python.framework import ops
  from tensorflow.python.ops.resource_variable_ops import ResourceVariable
  from tensorflow.core.protobuf import config_pb2
  from autodist.kernel.common.variable_utils import get_read_var_tensor
  from autodist.kernel.common.utils import replica_prefix
  from autodist.utils import logging
label_desc: Use image node_id 7 for calling the Remapper obj's underlying member method code with example usage: obj._set_default() and returns: final_fetches,  function, final_feeds, final_feeds