parameters: (tpu_assignment, mode,
                                     input_specs, get_next_ops)
return_statement: return input_specs, get_next_ops
docstring: Append core information to the set of dataset inputs.
docstring_summary: Append core information to the set of dataset inputs.
docstring_tokens: Append core information to the set of dataset inputs .
function: def _inject_tpu_inputs_for_dataset(tpu_assignment, mode,
                                     input_specs, get_next_ops):
    """Append core information to the set of dataset inputs."""
    # This is used during compilation to identify the current TPU core and enable
    # concatenation operations across cores.
    if mode not in [model_fn_lib.ModeKeys.TRAIN, model_fn_lib.ModeKeys.EVAL]:
      return input_specs, get_next_ops
  
    # Dataset inputs operate on per core basis.
    per_core_batch_size = input_specs[0].shape.as_list()[0]
  
    # Insert, at head, the tensor for core_id.
    assert len(get_next_ops) == tpu_assignment.num_towers
    for i in range(tpu_assignment.num_towers):
      core_id_constant = constant_op.constant(
          np.array([i] * per_core_batch_size).astype('int32'),
          dtype=dtypes.int32,
          name='cord_id_constant')
      get_next_ops[i] = [core_id_constant] + list(get_next_ops[i])
  
    # Insert the input spec at head also.
    input_specs = [tensor_spec.TensorSpec([per_core_batch_size], dtypes.int32)
                  ] + input_specs
  
    return input_specs, get_next_ops
function_tokens: def _inject_tpu_inputs_for_dataset ( tpu_assignment , mode , input_specs , get_next_ops ) : # This is used during compilation to identify the current TPU core and enable # concatenation operations across cores. if mode not in [ model_fn_lib . ModeKeys . TRAIN , model_fn_lib . ModeKeys . EVAL ] : return input_specs , get_next_ops # Dataset inputs operate on per core basis. per_core_batch_size = input_specs [ 0 ] . shape . as_list ( ) [ 0 ] # Insert, at head, the tensor for core_id. assert len ( get_next_ops ) == tpu_assignment . num_towers for i in range ( tpu_assignment . num_towers ) : core_id_constant = constant_op . constant ( np . array ( [ i ] * per_core_batch_size ) . astype ( 'int32' ) , dtype = dtypes . int32 , name = 'cord_id_constant' ) get_next_ops [ i ] = [ core_id_constant ] + list ( get_next_ops [ i ] ) # Insert the input spec at head also. input_specs = [ tensor_spec . TensorSpec ( [ per_core_batch_size ] , dtypes . int32 ) ] + input_specs return input_specs , get_next_ops