parameters: (computation,
                inputs=None,
                infeed_queue=None,
                global_tpu_id=None,
                name=None)
docstring: Builds a graph operator that runs a replicated TPU computation.
  
    Args:
      computation: a Python function that builds the computation to replicate.
      inputs: a list of lists of input tensors or None (equivalent to
        [[]]), indexed by [replica_num][input_num]. All replicas must
        have the same number of inputs.
      infeed_queue: if not None, the InfeedQueue from which to append a tuple
        of arguments as inputs to computation.
      global_tpu_id: if not None, a Numpy 2D array indicating the global
        id of each TPU device in the system. The outer dimension of the
        array is host task id, and the inner dimension is device ordinal,
        so e.g., global_tpu_id[x][y] indicates the global id of device
        /task:x/device:TPU_NODE:y.
      name: name of the operator.
    Returns:
      A list of lists of output tensors, indexed by [replica_num][output_num].
    Raises:
      ValueError: if all replicas do not have equal numbers of input tensors.
      ValueError: if the number of inputs per replica does not match
        the number of formal parameters to `computation`.
docstring_summary: Builds a graph operator that runs a replicated TPU computation.
docstring_tokens: Builds a graph operator that runs a replicated TPU computation .
function: def replicate(computation,
                inputs=None,
                infeed_queue=None,
                global_tpu_id=None,
                name=None):
    """Builds a graph operator that runs a replicated TPU computation.
  
    Args:
      computation: a Python function that builds the computation to replicate.
      inputs: a list of lists of input tensors or None (equivalent to
        [[]]), indexed by [replica_num][input_num]. All replicas must
        have the same number of inputs.
      infeed_queue: if not None, the InfeedQueue from which to append a tuple
        of arguments as inputs to computation.
      global_tpu_id: if not None, a Numpy 2D array indicating the global
        id of each TPU device in the system. The outer dimension of the
        array is host task id, and the inner dimension is device ordinal,
        so e.g., global_tpu_id[x][y] indicates the global id of device
        /task:x/device:TPU_NODE:y.
      name: name of the operator.
    Returns:
      A list of lists of output tensors, indexed by [replica_num][output_num].
    Raises:
      ValueError: if all replicas do not have equal numbers of input tensors.
      ValueError: if the number of inputs per replica does not match
        the number of formal parameters to `computation`.
    """
    if name is None:
      name = "TPUReplicate"
    inputs = [[]] if inputs is None else inputs
  
    if global_tpu_id is not None:
      # Turn the Numpy array into a flattened list.
      global_tpu_id = global_tpu_id.flatten().tolist()
  
    if ((not isinstance(inputs, list)) or
        any(not isinstance(inp, (list, tuple)) for inp in inputs)):
      raise TypeError("tpu.replicate() inputs must be a list of lists/tuples")
  
    num_replicas = len(inputs)
  
    # No replicas? Nothing to do.
    if num_replicas == 0:
      return []
  
    # Converts inputs to Tensors.
    inputs = [[ops.convert_to_tensor(x) for x in inp] for inp in inputs]
  
    # Verifies that all replicas have matching numbers and types of inputs
    input_types = [x.dtype for x in inputs[0]]
    input_arity = len(input_types)
    for i in range(num_replicas):
      if len(inputs[i]) != input_arity:
        raise ValueError("Replicas must have the same number of inputs. "
                         "Replica 0 had {} inputs, replica {} had {} "
                         "inputs.".format(input_arity, i, len(inputs[i])))
  
      types = [x.dtype for x in inputs[i]]
      if types != input_types:
        raise ValueError(
            "Replicas must have matching input types. Replica 0 had "
            "input types {}, replica {} had input types {}".format(
                input_types, i, types))
  
    arg_error = tpu_function.check_function_argument_count(
        computation, input_arity, infeed_queue)
    if arg_error is not None:
      if infeed_queue is None:
        raise TypeError(
            "Supplied computation cannot be called with the specified inputs. "
            "You specified %d inputs: %s, but the computation needs %s" % (
                input_arity, str([i.name for i in inputs[0]]), arg_error))
      else:
        raise TypeError(
            "Supplied computation cannot be called with the specified inputs. "
            "You specified %d inputs: %s and %d additional inputs from infeed,"
            " but the computation needs %s" % (input_arity, str(
                [i.name
                 for i in inputs[0]]), infeed_queue.number_of_tuple_elements,
                                               arg_error))
  
    graph = ops.get_default_graph()
  
    with ops.name_scope(name, "replicate"):
      # Fan-in: Builds a TPUReplicatedInput node for each input.
      computation_inputs = []
      for i in range(0, input_arity):
        replicas = [inputs[replica][i] for replica in xrange(num_replicas)]
        computation_inputs.append(
            tpu_ops.tpu_replicated_input(replicas, name="input{}".format(i)))
  
      context = TPUReplicateContext(name=graph.unique_name("cluster"))
      try:
        context.Enter()
  
        metadata = tpu_ops.tpu_replicate_metadata(
            num_replicas=num_replicas, global_tpu_id=global_tpu_id)
  
        with tpu_function.tpu_shard_context(
            num_replicas), ops.control_dependencies([metadata]):
  
          # The EncapsulateTPUComputations rewrite needs to identify the
          # replicated arguments inside each computation. Adds identity operators
          # tagged with an attribute _tpu_replicated_input to identify the
          # replicated inputs.
          # pylint: disable=protected-access
          with graph._attr_scope({"_tpu_replicated_input":
                                  attr_value_pb2.AttrValue(b=True)}):
            computation_inputs = [
                array_ops.identity(x, name="replicated_input_{}".format(i))
                for i, x in enumerate(computation_inputs)]
          # pylint: enable=protected-access
  
          # If there is an infeed queue, adds the dequeued values to the
          # computation's inputs.
          if infeed_queue is not None:
            infeed_queue.set_number_of_shards(num_replicas)
            for t in infeed_queue.generate_dequeue_op():
              computation_inputs.append(t)
  
          # Only resource variables work inside a TPU computation, so turn on
          # resource variables for the computation.
          # TODO(phawkins): consider removing this code. It will
          # be less confusing to clients if they knowingly choose to use resource
          # variables.
          vscope = variable_scope.get_variable_scope()
          saved_use_resource = vscope.use_resource
          vscope.set_use_resource(True)
  
          outputs = computation(*computation_inputs)
  
          vscope.set_use_resource(saved_use_resource)
  
        # If the computation only returned one value, makes it a tuple.
        if not isinstance(outputs, (list, tuple)):
          outputs = (outputs,)
  
        try:
          with ops.device(core(0)):
            outputs = [
                o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o)
                for o in outputs
            ]
        except Exception as e:
          raise ValueError(
              "TPU function return values must all either be Operations or "
              "convertible to Tensors. Got '%s'" % str(e))
  
        # Separates the returned Operations and Tensors.
        output_operations = [o for o in outputs if isinstance(o, ops.Operation)]
        output_tensors = [o for o in outputs
                          if not isinstance(o, ops.Operation)]
  
        if outputs != output_tensors + output_operations:
          raise ValueError(
              "TPU functions must return zero-or more Tensor values followed by "
              "zero or more Operations.")
        output_arity = len(output_tensors)
  
        # Wraps outputs in Identity ops. Otherwise a replicated input copied
        # straight to an output would bypass the replicate(). This would be bad
        # because the TPUReplicatedInput/TPUReplicatedOutput operator would not
        # be rewritten away, leading to a runtime error.
        # TODO(phawkins): extend the rewrite to elide these nodes instead.
        with ops.device(core(0)):
          output_tensors = [array_ops.identity(x) for x in output_tensors]
      finally:
        context.Exit()
  
      # Fan-out: Builds a TPUReplicatedOutput node for each output.
      outputs = [tpu_ops.tpu_replicated_output(output_tensors[i], num_replicas,
                                               name="output{}".format(i))
                 for i in xrange(output_arity)]
  
      with ops.control_dependencies(output_operations):
        if output_arity == 0:
          # Returns a list of NoOps dependent on the replication Op, indexed by
          # [replica_num].
          return [
              control_flow_ops.no_op(name="%s_shard_%d" % (name, i))
              for i in range(num_replicas)
          ]
        else:
          # Wraps the outputs in identity operators so the names of any possible
          # `fetch` nodes are preserved by the replication rewrite.
          return [
              [array_ops.identity(outputs[out][replica],
                                  name="output_%d_shard_%d" % (out, replica))
               for out in xrange(output_arity)]
              for replica in xrange(num_replicas)
          ]
function_tokens: def replicate ( computation , inputs = None , infeed_queue = None , global_tpu_id = None , name = None ) : if name is None : name = "TPUReplicate" inputs = [ [ ] ] if inputs is None else inputs if global_tpu_id is not None : # Turn the Numpy array into a flattened list. global_tpu_id = global_tpu_id . flatten ( ) . tolist ( ) if ( ( not isinstance ( inputs , list ) ) or any ( not isinstance ( inp , ( list , tuple ) ) for inp in inputs ) ) : raise TypeError ( "tpu.replicate() inputs must be a list of lists/tuples" ) num_replicas = len ( inputs ) # No replicas? Nothing to do. if num_replicas == 0 : return [ ] # Converts inputs to Tensors. inputs = [ [ ops . convert_to_tensor ( x ) for x in inp ] for inp in inputs ] # Verifies that all replicas have matching numbers and types of inputs input_types = [ x . dtype for x in inputs [ 0 ] ] input_arity = len ( input_types ) for i in range ( num_replicas ) : if len ( inputs [ i ] ) != input_arity : raise ValueError ( "Replicas must have the same number of inputs. " "Replica 0 had {} inputs, replica {} had {} " "inputs." . format ( input_arity , i , len ( inputs [ i ] ) ) ) types = [ x . dtype for x in inputs [ i ] ] if types != input_types : raise ValueError ( "Replicas must have matching input types. Replica 0 had " "input types {}, replica {} had input types {}" . format ( input_types , i , types ) ) arg_error = tpu_function . check_function_argument_count ( computation , input_arity , infeed_queue ) if arg_error is not None : if infeed_queue is None : raise TypeError ( "Supplied computation cannot be called with the specified inputs. " "You specified %d inputs: %s, but the computation needs %s" % ( input_arity , str ( [ i . name for i in inputs [ 0 ] ] ) , arg_error ) ) else : raise TypeError ( "Supplied computation cannot be called with the specified inputs. " "You specified %d inputs: %s and %d additional inputs from infeed," " but the computation needs %s" % ( input_arity , str ( [ i . name for i in inputs [ 0 ] ] ) , infeed_queue . number_of_tuple_elements , arg_error ) ) graph = ops . get_default_graph ( ) with ops . name_scope ( name , "replicate" ) : # Fan-in: Builds a TPUReplicatedInput node for each input. computation_inputs = [ ] for i in range ( 0 , input_arity ) : replicas = [ inputs [ replica ] [ i ] for replica in xrange ( num_replicas ) ] computation_inputs . append ( tpu_ops . tpu_replicated_input ( replicas , name = "input{}" . format ( i ) ) ) context = TPUReplicateContext ( name = graph . unique_name ( "cluster" ) ) try : context . Enter ( ) metadata = tpu_ops . tpu_replicate_metadata ( num_replicas = num_replicas , global_tpu_id = global_tpu_id ) with tpu_function . tpu_shard_context ( num_replicas ) , ops . control_dependencies ( [ metadata ] ) : # The EncapsulateTPUComputations rewrite needs to identify the # replicated arguments inside each computation. Adds identity operators # tagged with an attribute _tpu_replicated_input to identify the # replicated inputs. # pylint: disable=protected-access with graph . _attr_scope ( { "_tpu_replicated_input" : attr_value_pb2 . AttrValue ( b = True ) } ) : computation_inputs = [ array_ops . identity ( x , name = "replicated_input_{}" . format ( i ) ) for i , x in enumerate ( computation_inputs ) ] # pylint: enable=protected-access # If there is an infeed queue, adds the dequeued values to the # computation's inputs. if infeed_queue is not None : infeed_queue . set_number_of_shards ( num_replicas ) for t in infeed_queue . generate_dequeue_op ( ) : computation_inputs . append ( t ) # Only resource variables work inside a TPU computation, so turn on # resource variables for the computation. # TODO(phawkins): consider removing this code. It will # be less confusing to clients if they knowingly choose to use resource # variables. vscope = variable_scope . get_variable_scope ( ) saved_use_resource = vscope . use_resource vscope . set_use_resource ( True ) outputs = computation ( * computation_inputs ) vscope . set_use_resource ( saved_use_resource ) # If the computation only returned one value, makes it a tuple. if not isinstance ( outputs , ( list , tuple ) ) : outputs = ( outputs , ) try : with ops . device ( core ( 0 ) ) : outputs = [ o if isinstance ( o , ops . Operation ) else ops . convert_to_tensor ( o ) for o in outputs ] except Exception as e : raise ValueError ( "TPU function return values must all either be Operations or " "convertible to Tensors. Got '%s'" % str ( e ) ) # Separates the returned Operations and Tensors. output_operations = [ o for o in outputs if isinstance ( o , ops . Operation ) ] output_tensors = [ o for o in outputs if not isinstance ( o , ops . Operation ) ] if outputs != output_tensors + output_operations : raise ValueError ( "TPU functions must return zero-or more Tensor values followed by " "zero or more Operations." ) output_arity = len ( output_tensors ) # Wraps outputs in Identity ops. Otherwise a replicated input copied # straight to an output would bypass the replicate(). This would be bad # because the TPUReplicatedInput/TPUReplicatedOutput operator would not # be rewritten away, leading to a runtime error. # TODO(phawkins): extend the rewrite to elide these nodes instead. with ops . device ( core ( 0 ) ) : output_tensors = [ array_ops . identity ( x ) for x in output_tensors ] finally : context . Exit ( ) # Fan-out: Builds a TPUReplicatedOutput node for each output. outputs = [ tpu_ops . tpu_replicated_output ( output_tensors [ i ] , num_replicas , name = "output{}" . format ( i ) ) for i in xrange ( output_arity ) ] with ops . control_dependencies ( output_operations ) : if output_arity == 0 : # Returns a list of NoOps dependent on the replication Op, indexed by # [replica_num]. return [ control_flow_ops . no_op ( name = "%s_shard_%d" % ( name , i ) ) for i in range ( num_replicas ) ] else : # Wraps the outputs in identity operators so the names of any possible # `fetch` nodes are preserved by the replication rewrite. return [ [ array_ops . identity ( outputs [ out ] [ replica ] , name = "output_%d_shard_%d" % ( out , replica ) ) for out in xrange ( output_arity ) ] for replica in xrange ( num_replicas ) ]