Gathers slices from `params` according to `indices` with batch dims. This operation is similar to `gather`, but it assumes that the leading `N` dimensions of `indices` and `params` are batch dimensions, and performs a gather within each batch. In particular, when using this operation with `N
(params, indices, name=None)
| 31 | # ragged.batch_gather |
| 32 | #=============================================================================== |
| 33 | def batch_gather(params, indices, name=None): |
| 34 | """Gathers slices from `params` according to `indices` with batch dims. |
| 35 | |
| 36 | This operation is similar to `gather`, but it assumes that the leading `N` |
| 37 | dimensions of `indices` and `params` are batch dimensions, and performs a |
| 38 | gather within each batch. In particular, when using this operation with `N` |
| 39 | batch dimensions `B1...BN`: |
| 40 | |
| 41 | * `indices` has shape `[B1...BN, I]` |
| 42 | * `params` has shape `[B1...BN, P1...PM]`. |
| 43 | * `result` has shape `[B1...BN, I, P2...PM]`. |
| 44 | * `result[b1...bN, i, p2...pM] = |
| 45 | params[b1...bN, indices[b1...bN, i], p2...pM]` |
| 46 | |
| 47 | Args: |
| 48 | params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`, |
| 49 | `M>0`). |
| 50 | indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`). |
| 51 | name: A name for the operation (optional). |
| 52 | |
| 53 | Returns: |
| 54 | A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`. |
| 55 | `result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`. |
| 56 | |
| 57 | #### Example: |
| 58 | ```python |
| 59 | >>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']]) |
| 60 | >>> indices = tf.ragged.constant([[1, 2, 0], [], [], [0, 0]]) |
| 61 | >>> tf.compat.v1.batch_gather(params, indices) |
| 62 | [['b', 'c', 'a'], [], [], ['e', 'e']] |
| 63 | ``` |
| 64 | """ |
| 65 | if not (ragged_tensor.is_ragged(params) or ragged_tensor.is_ragged(indices)): |
| 66 | return array_ops.batch_gather(params, indices, name) |
| 67 | |
| 68 | with ops.name_scope(name, 'RaggedBatchGather', [params, indices]): |
| 69 | params = ragged_tensor.convert_to_tensor_or_ragged_tensor( |
| 70 | params, name='params') |
| 71 | indices = ragged_tensor.convert_to_tensor_or_ragged_tensor( |
| 72 | indices, name='indices') |
| 73 | params, indices = ragged_tensor.match_row_splits_dtypes(params, indices) |
| 74 | indices_ndims = indices.shape.ndims |
| 75 | if indices_ndims is None: |
| 76 | raise ValueError( |
| 77 | 'batch_gather does not allow indices with unknown shape.') |
| 78 | if indices_ndims == 0: |
| 79 | raise ValueError('indices.rank must be at least 1.') |
| 80 | |
| 81 | if ragged_tensor.is_ragged(indices): |
| 82 | # If the outermost ragged dimension is a batch dimension, recurse. |
| 83 | if indices_ndims > 2: |
| 84 | if not ragged_tensor.is_ragged(params): |
| 85 | raise ValueError('batch shape from indices does ' |
| 86 | 'not match params shape') |
| 87 | checks = [check_ops.assert_equal(params.row_splits, indices.row_splits)] |
| 88 | with ops.control_dependencies(checks): |
| 89 | return ragged_tensor.RaggedTensor.from_row_splits( |
| 90 | batch_gather(params.values, indices.values), indices.row_splits, |
nothing calls this directly
no test coverage detected