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`, where `params` is interpreted as a partitioning of a large embedding tensor. `params` may be a `PartitionedVariable`
(
params,
ids,
partition_strategy="mod",
name=None,
validate_indices=True, # pylint: disable=unused-argument
max_norm=None,
ev_init_value=None,
blocknums=None,
counts=None)
| 363 | |
| 364 | @tf_export(v1=["nn.embedding_lookup"]) |
| 365 | def embedding_lookup( |
| 366 | params, |
| 367 | ids, |
| 368 | partition_strategy="mod", |
| 369 | name=None, |
| 370 | validate_indices=True, # pylint: disable=unused-argument |
| 371 | max_norm=None, |
| 372 | ev_init_value=None, |
| 373 | blocknums=None, |
| 374 | counts=None): |
| 375 | """Looks up `ids` in a list of embedding tensors. |
| 376 | This function is used to perform parallel lookups on the list of |
| 377 | tensors in `params`. It is a generalization of |
| 378 | `tf.gather`, where `params` is |
| 379 | interpreted as a partitioning of a large embedding tensor. `params` may be |
| 380 | a `PartitionedVariable` as returned by using `tf.compat.v1.get_variable()` |
| 381 | with a |
| 382 | partitioner. |
| 383 | If `len(params) > 1`, each element `id` of `ids` is partitioned between |
| 384 | the elements of `params` according to the `partition_strategy`. |
| 385 | In all strategies, if the id space does not evenly divide the number of |
| 386 | partitions, each of the first `(max_id + 1) % len(params)` partitions will |
| 387 | be assigned one more id. |
| 388 | If `partition_strategy` is `"mod"`, we assign each id to partition |
| 389 | `p = id % len(params)`. For instance, |
| 390 | 13 ids are split across 5 partitions as: |
| 391 | `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]` |
| 392 | If `partition_strategy` is `"div"`, we assign ids to partitions in a |
| 393 | contiguous manner. In this case, 13 ids are split across 5 partitions as: |
| 394 | `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]` |
| 395 | The results of the lookup are concatenated into a dense |
| 396 | tensor. The returned tensor has shape `shape(ids) + shape(params)[1:]`. |
| 397 | Args: |
| 398 | params: A single tensor representing the complete embedding tensor, or a |
| 399 | list of P tensors all of same shape except for the first dimension, |
| 400 | representing sharded embedding tensors. Alternatively, a |
| 401 | `PartitionedVariable`, created by partitioning along dimension 0. Each |
| 402 | element must be appropriately sized for the given `partition_strategy`. |
| 403 | ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked |
| 404 | up in `params`. |
| 405 | partition_strategy: A string specifying the partitioning strategy, relevant |
| 406 | if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default |
| 407 | is `"mod"`. |
| 408 | name: A name for the operation (optional). |
| 409 | validate_indices: DEPRECATED. If this operation is assigned to CPU, values |
| 410 | in `indices` are always validated to be within range. If assigned to GPU, |
| 411 | out-of-bound indices result in safe but unspecified behavior, which may |
| 412 | include raising an error. |
| 413 | max_norm: If not `None`, each embedding is clipped if its l2-norm is larger |
| 414 | than this value. |
| 415 | Returns: |
| 416 | A `Tensor` with the same type as the tensors in `params`. |
| 417 | Raises: |
| 418 | ValueError: If `params` is empty. |
| 419 | """ |
| 420 | return _embedding_lookup_and_transform( |
| 421 | params=params, |
| 422 | ids=ids, |
no test coverage detected