Remapper

Feed and Fetch Remapper.

class Remapper(graph_transformer, graph_item)

Bases: object

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.

remap_callable_options(callable_options)

Remap Callable Options.

Parameters

callable_options – A CallableOptions protocol buffer message describing the computation that will be performed by the callable.

Returns

A new CallableOptions

as_default()

Ensure the current one as the default remapper; otherwise switch to it.