Gathers ragged slices from `params` axis `0` according to `indices`. Returns `RaggedTensor` output, such that: ```python output.shape = indices.shape + params.shape[1:] output.ragged_rank = indices.shape.ndims + params.ragged_rank output[i...j, d0...dn] = params[indices[i...j], d0...dn]
(params, indices, validate_indices=None, axis=0, batch_dims=0,
name=None)
| 35 | #=============================================================================== |
| 36 | # TODO(edloper): Add an `axis` argument |
| 37 | def gather(params, indices, validate_indices=None, axis=0, batch_dims=0, |
| 38 | name=None): |
| 39 | """Gathers ragged slices from `params` axis `0` according to `indices`. |
| 40 | |
| 41 | Returns `RaggedTensor` output, such that: |
| 42 | |
| 43 | ```python |
| 44 | output.shape = indices.shape + params.shape[1:] |
| 45 | output.ragged_rank = indices.shape.ndims + params.ragged_rank |
| 46 | output[i...j, d0...dn] = params[indices[i...j], d0...dn] |
| 47 | ``` |
| 48 | |
| 49 | `params` may be ragged. `indices` may be ragged. |
| 50 | `indices` must have dtype `int32` or `int64`. If any index is out of bounds, |
| 51 | then an error is returned. |
| 52 | |
| 53 | Examples: |
| 54 | |
| 55 | ```python |
| 56 | >>> params = tf.constant(['a', 'b', 'c', 'd', 'e']) |
| 57 | >>> indices = tf.constant([3, 1, 2, 1, 0]) |
| 58 | >>> ragged_params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']]) |
| 59 | >>> ragged_indices = tf.ragged.constant([[3, 1, 2], [1], [], [0]]) |
| 60 | |
| 61 | >>> print ragged.gather(params, ragged_indices) |
| 62 | [['d', 'b', 'c'], ['b'], [], ['a']] |
| 63 | |
| 64 | >>> print ragged.gather(ragged_params, indices) |
| 65 | [['e'], ['d'], [], ['d'], ['a', 'b', 'c']] |
| 66 | |
| 67 | >>> print ragged.gather(ragged_params, ragged_indices) |
| 68 | [[['e'], ['d'], []], [['d']], [], [['a', 'b', 'c']]] |
| 69 | ``` |
| 70 | |
| 71 | Args: |
| 72 | params: The potentially ragged tensor from which to gather values. Must be |
| 73 | at least rank 1. |
| 74 | indices: The potentially ragged tensor indicating which values to gather. |
| 75 | Must have dtype `int32` or `int64`. Values must be in the range `[0, |
| 76 | params.shape[0]]`. |
| 77 | validate_indices: Ignored. |
| 78 | axis: Must be zero. |
| 79 | batch_dims: Must be zero. |
| 80 | name: A name for the operation (optional). |
| 81 | |
| 82 | Returns: |
| 83 | A `RaggedTensor`, where `output.dtype=params.dtype` and |
| 84 | `output.shape=indices.shape + params.shape[1:]` and |
| 85 | `output.ragged_rank=indices.shape.ndims + params.ragged_rank`. |
| 86 | |
| 87 | Raises: |
| 88 | ValueError: If indices.shape.ndims is not known statically. |
| 89 | """ |
| 90 | del validate_indices |
| 91 | if not isinstance(axis, int) or axis != 0: |
| 92 | raise ValueError('axis != 0 is not supported for ragged gather yet.') |
| 93 | if not isinstance(batch_dims, int) or batch_dims != 0: |
| 94 | raise ValueError('batch_dims != 0 is not supported for ragged gather yet.') |
no test coverage detected