parameters: (params, ids, partition_strategy="mod", name=None,
                       validate_indices=True)
docstring: Looks up `ids` in a list of embedding tensors.
  
    This function is used to perform parallel lookups on the list of
    tensors in `params`.  It is a generalization of
    [`tf.gather()`](../../api_docs/python/array_ops.md#gather), where `params` is
    interpreted as a partitioning of a large embedding tensor.  `params` may be
    a `PartitionedVariable` as returned by using `tf.get_variable()` with a
    partitioner.
  
    If `len(params) > 1`, each element `id` of `ids` is partitioned between
    the elements of `params` according to the `partition_strategy`.
    In all strategies, if the id space does not evenly divide the number of
    partitions, each of the first `(max_id + 1) % len(params)` partitions will
    be assigned one more id.
  
    If `partition_strategy` is `"mod"`, we assign each id to partition
    `p = id % len(params)`. For instance,
    13 ids are split across 5 partitions as:
    `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`
  
    If `partition_strategy` is `"div"`, we assign ids to partitions in a
    contiguous manner. In this case, 13 ids are split across 5 partitions as:
    `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`
  
    The results of the lookup are concatenated into a dense
    tensor. The returned tensor has shape `shape(ids) + shape(params)[1:]`.
  
    Args:
      params: A list of tensors with the same type and which can be concatenated
        along dimension 0. Alternatively, a `PartitionedVariable`, created by
        partitioning along dimension 0.  Each element must be appropriately sized
        for the given `partition_strategy`.
      ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked
        up in `params`.
      partition_strategy: A string specifying the partitioning strategy, relevant
        if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
        is `"mod"`.
      name: A name for the operation (optional).
      validate_indices: Whether or not to validate gather indices.
  
    Returns:
      A `Tensor` with the same type as the tensors in `params`.
  
    Raises:
      ValueError: If `params` is empty.
docstring_summary: Looks up `ids` in a list of embedding tensors.
docstring_tokens: Looks up ids in a list of embedding tensors .
function: def embedding_lookup(params, ids, partition_strategy="mod", name=None,
                       validate_indices=True):
    """Looks up `ids` in a list of embedding tensors.
  
    This function is used to perform parallel lookups on the list of
    tensors in `params`.  It is a generalization of
    [`tf.gather()`](../../api_docs/python/array_ops.md#gather), where `params` is
    interpreted as a partitioning of a large embedding tensor.  `params` may be
    a `PartitionedVariable` as returned by using `tf.get_variable()` with a
    partitioner.
  
    If `len(params) > 1`, each element `id` of `ids` is partitioned between
    the elements of `params` according to the `partition_strategy`.
    In all strategies, if the id space does not evenly divide the number of
    partitions, each of the first `(max_id + 1) % len(params)` partitions will
    be assigned one more id.
  
    If `partition_strategy` is `"mod"`, we assign each id to partition
    `p = id % len(params)`. For instance,
    13 ids are split across 5 partitions as:
    `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`
  
    If `partition_strategy` is `"div"`, we assign ids to partitions in a
    contiguous manner. In this case, 13 ids are split across 5 partitions as:
    `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`
  
    The results of the lookup are concatenated into a dense
    tensor. The returned tensor has shape `shape(ids) + shape(params)[1:]`.
  
    Args:
      params: A list of tensors with the same type and which can be concatenated
        along dimension 0. Alternatively, a `PartitionedVariable`, created by
        partitioning along dimension 0.  Each element must be appropriately sized
        for the given `partition_strategy`.
      ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked
        up in `params`.
      partition_strategy: A string specifying the partitioning strategy, relevant
        if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
        is `"mod"`.
      name: A name for the operation (optional).
      validate_indices: Whether or not to validate gather indices.
  
    Returns:
      A `Tensor` with the same type as the tensors in `params`.
  
    Raises:
      ValueError: If `params` is empty.
    """
    if params is None or params == []:  # pylint: disable=g-explicit-bool-comparison
      raise ValueError("Need at least one param")
    if isinstance(params, variables.PartitionedVariable):
      params = list(params)  # Iterate to get the underlying Variables.
    if not isinstance(params, list):
      params = [params]
    with ops.name_scope(name, "embedding_lookup", params + [ids]) as name:
      np = len(params)  # Number of partitions
      params = ops.convert_n_to_tensor_or_indexed_slices(params, name="params")
      if np == 1:
        with ops.colocate_with(params[0]):
          return array_ops.gather(params[0], ids, name=name,
                                  validate_indices=validate_indices)
      else:
        ids = ops.convert_to_tensor(ids, name="ids")
        flat_ids = array_ops.reshape(ids, [-1])
        original_indices = math_ops.range(array_ops.size(flat_ids))
  
        # Create p_assignments and set new_ids depending on the strategy.
        if partition_strategy == "mod":
          p_assignments = flat_ids % np
          new_ids = flat_ids // np
        elif partition_strategy == "div":
          # Compute num_total_ids as the sum of dim-0 of params, then assign to
          # partitions based on a constant number of ids per partition. Optimize
          # if we already know the full shape statically.
          dim_0_size = params[0].get_shape()[0]
          for p in xrange(1, np):
            dim_0_size += params[p].get_shape()[0]
          if dim_0_size.value:
            num_total_ids = constant_op.constant(dim_0_size.value, flat_ids.dtype)
          else:
            dim_0_sizes = []
            for p in xrange(np):
              if params[p].get_shape()[0].value is not None:
                dim_0_sizes.append(params[p].get_shape()[0].value)
              else:
                with ops.colocate_with(params[p]):
                  dim_0_sizes.append(array_ops.shape(params[p])[0])
            num_total_ids = math_ops.reduce_sum(
                math_ops.cast(array_ops.pack(dim_0_sizes), flat_ids.dtype))
          ids_per_partition = num_total_ids // np
          extras = num_total_ids % np
  
          p_assignments = math_ops.maximum(
              flat_ids // (ids_per_partition + 1),
              (flat_ids - extras) // ids_per_partition)
  
          # Emulate a conditional using a boolean indicator tensor
          is_in_first_extras_partitions = math_ops.cast(
              p_assignments < extras, flat_ids.dtype)
          new_ids = (
              is_in_first_extras_partitions * (
                  flat_ids % (ids_per_partition + 1)) +
              (1 - is_in_first_extras_partitions) * (
                  (flat_ids - extras) % ids_per_partition))
        else:
          raise ValueError("Unrecognized partition strategy: " +
                           partition_strategy)
  
        # Cast partition assignments to int32 for use in dynamic_partition.
        # There really should not be more than 2^32 partitions.
        p_assignments = math_ops.cast(p_assignments, dtypes.int32)
        # Partition list of ids based on assignments into np separate lists
        gather_ids = data_flow_ops.dynamic_partition(new_ids, p_assignments, np)
        # Similarly, partition the original indices.
        pindices = data_flow_ops.dynamic_partition(original_indices,
                                                   p_assignments, np)
        # Do np separate lookups, finding embeddings for plist[p] in params[p]
        partitioned_result = []
        for p in xrange(np):
          with ops.colocate_with(params[p]):
            partitioned_result.append(array_ops.gather(
                params[p], gather_ids[p],
                validate_indices=validate_indices))
        # Stitch these back together
        ret = data_flow_ops.dynamic_stitch(pindices, partitioned_result,
                                           name=name)
        # Reshape to reverse the flattening of ids.
        element_shape = params[0].get_shape()[1:]
        for p in params[1:]:
          element_shape = element_shape.merge_with(p.get_shape()[1:])
        if element_shape.is_fully_defined():
          ret = array_ops.reshape(ret, array_ops.concat(0, [
              array_ops.shape(ids), element_shape]))
        else:
          # It's important that we compute params[0].shape on the right device
          # to avoid data motion.
          with ops.colocate_with(params[0]):
            params_shape = array_ops.shape(params[0])
          ret = array_ops.reshape(ret, array_ops.concat(0, [
              array_ops.shape(ids), array_ops.slice(params_shape, [1], [-1])]))
        # output shape = ids.shape + params[*].shape[1:]
        # Normally the reshape is sufficient, but setting shape explicitly
        # teaches shape inference that params[1:].get_shape() matters.
        ret.set_shape(ids.get_shape().concatenate(element_shape))
        return ret
function_tokens: def embedding_lookup ( params , ids , partition_strategy = "mod" , name = None , validate_indices = True ) : if params is None or params == [ ] : # pylint: disable=g-explicit-bool-comparison raise ValueError ( "Need at least one param" ) if isinstance ( params , variables . PartitionedVariable ) : params = list ( params ) # Iterate to get the underlying Variables. if not isinstance ( params , list ) : params = [ params ] with ops . name_scope ( name , "embedding_lookup" , params + [ ids ] ) as name : np = len ( params ) # Number of partitions params = ops . convert_n_to_tensor_or_indexed_slices ( params , name = "params" ) if np == 1 : with ops . colocate_with ( params [ 0 ] ) : return array_ops . gather ( params [ 0 ] , ids , name = name , validate_indices = validate_indices ) else : ids = ops . convert_to_tensor ( ids , name = "ids" ) flat_ids = array_ops . reshape ( ids , [ - 1 ] ) original_indices = math_ops . range ( array_ops . size ( flat_ids ) ) # Create p_assignments and set new_ids depending on the strategy. if partition_strategy == "mod" : p_assignments = flat_ids % np new_ids = flat_ids // np elif partition_strategy == "div" : # Compute num_total_ids as the sum of dim-0 of params, then assign to # partitions based on a constant number of ids per partition. Optimize # if we already know the full shape statically. dim_0_size = params [ 0 ] . get_shape ( ) [ 0 ] for p in xrange ( 1 , np ) : dim_0_size += params [ p ] . get_shape ( ) [ 0 ] if dim_0_size . value : num_total_ids = constant_op . constant ( dim_0_size . value , flat_ids . dtype ) else : dim_0_sizes = [ ] for p in xrange ( np ) : if params [ p ] . get_shape ( ) [ 0 ] . value is not None : dim_0_sizes . append ( params [ p ] . get_shape ( ) [ 0 ] . value ) else : with ops . colocate_with ( params [ p ] ) : dim_0_sizes . append ( array_ops . shape ( params [ p ] ) [ 0 ] ) num_total_ids = math_ops . reduce_sum ( math_ops . cast ( array_ops . pack ( dim_0_sizes ) , flat_ids . dtype ) ) ids_per_partition = num_total_ids // np extras = num_total_ids % np p_assignments = math_ops . maximum ( flat_ids // ( ids_per_partition + 1 ) , ( flat_ids - extras ) // ids_per_partition ) # Emulate a conditional using a boolean indicator tensor is_in_first_extras_partitions = math_ops . cast ( p_assignments < extras , flat_ids . dtype ) new_ids = ( is_in_first_extras_partitions * ( flat_ids % ( ids_per_partition + 1 ) ) + ( 1 - is_in_first_extras_partitions ) * ( ( flat_ids - extras ) % ids_per_partition ) ) else : raise ValueError ( "Unrecognized partition strategy: " + partition_strategy ) # Cast partition assignments to int32 for use in dynamic_partition. # There really should not be more than 2^32 partitions. p_assignments = math_ops . cast ( p_assignments , dtypes . int32 ) # Partition list of ids based on assignments into np separate lists gather_ids = data_flow_ops . dynamic_partition ( new_ids , p_assignments , np ) # Similarly, partition the original indices. pindices = data_flow_ops . dynamic_partition ( original_indices , p_assignments , np ) # Do np separate lookups, finding embeddings for plist[p] in params[p] partitioned_result = [ ] for p in xrange ( np ) : with ops . colocate_with ( params [ p ] ) : partitioned_result . append ( array_ops . gather ( params [ p ] , gather_ids [ p ] , validate_indices = validate_indices ) ) # Stitch these back together ret = data_flow_ops . dynamic_stitch ( pindices , partitioned_result , name = name ) # Reshape to reverse the flattening of ids. element_shape = params [ 0 ] . get_shape ( ) [ 1 : ] for p in params [ 1 : ] : element_shape = element_shape . merge_with ( p . get_shape ( ) [ 1 : ] ) if element_shape . is_fully_defined ( ) : ret = array_ops . reshape ( ret , array_ops . concat ( 0 , [ array_ops . shape ( ids ) , element_shape ] ) ) else : # It's important that we compute params[0].shape on the right device # to avoid data motion. with ops . colocate_with ( params [ 0 ] ) : params_shape = array_ops . shape ( params [ 0 ] ) ret = array_ops . reshape ( ret , array_ops . concat ( 0 , [ array_ops . shape ( ids ) , array_ops . slice ( params_shape , [ 1 ] , [ - 1 ] ) ] ) ) # output shape = ids.shape + params[*].shape[1:] # Normally the reshape is sufficient, but setting shape explicitly # teaches shape inference that params[1:].get_shape() matters. ret . set_shape ( ids . get_shape ( ) . concatenate ( element_shape ) ) return ret