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: Remap the feeds to the right element in the transformed graph.
  
  For example, there are N copies of a placeholder for N replicas
    and we have to feed all of them with tensors.
  
  Args:
      feed: feed graph element or name
      feed_val: feed value
  
  Returns:
      List of (new_feed, new_feed_value) pairs
docstr_tok: ["Remap","the","feeds","to","the","right","element","in","the","transformed","graph",".","For","example",",","there","are","N","copies","of","a","placeholder","for","N","replicas","and","we","have","to","feed","all","of","them","with","tensors",".","Args",":","feed",":","feed","graph","element","or","name","feed_val",":","feed","value","Returns",":","List","of","(","new_feed",",","new_feed_value",")","pairs"]
code: def _remap_feed(self, feed, feed_val=None):
      """
      Remap the feeds to the right element in the transformed graph.
  
      For example, there are N copies of a placeholder for N replicas
        and we have to feed all of them with tensors.
  
      Args:
          feed: feed graph element or name
          feed_val: feed value
  
      Returns:
          List of (new_feed, new_feed_value) pairs
      """
      feed_name = feed if isinstance(feed, str) else feed.name
      try:
          transformed_feeds = [
              self._graph_item.graph.as_graph_element(feed_name)
          ]
      except KeyError:
          transformed_feeds = [
              self._graph_item.graph.as_graph_element(
                  ops.prepend_name_scope(feed_name, replica_prefix(i))
              )
              for i in range(self._graph_transformer.num_local_replicas)
          ]
  
      num_replicated_feeds = self._graph_transformer.num_local_replicas
      feed = feed if not isinstance(feed, str) else transformed_feeds[0]
  
      def expand_feed_val(feed_val, feed=feed):
          """Given a original feed or replicated feed, expand the feed value."""
          # If we have replicated placeholders with undefined (polymorphic) shape, we split the feed_val across it;
          #  otherwise we feed all replicated placeholders the same feed_val
          polymorphic_dim = self._polymorphic_dim(feed)
          if polymorphic_dim:
              feed_vals = np.array_split(
                  np.asarray(feed_val),
                  num_replicated_feeds,
                  axis=polymorphic_dim,
              )
          else:
              feed_vals = [
                  feed_val for _ in range(num_replicated_feeds)
              ]
          return feed_vals
  
      if feed_val is not None:
          feed_vals = expand_feed_val(feed_val)
          transformed_feeds = list(zip(transformed_feeds, feed_vals))
      return transformed_feeds, expand_feed_val
  
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 4 for calling the Remapper obj's underlying member method code with example usage: obj._remap_feed(feed, feed_val) and returns: transformed_feeds,  expand_feed_val, feed_vals