Helper for indexing and slicing ragged tensors with __getitem__(). Extracts the specified piece of the `rt_input`. See `RaggedTensor.__getitem__` for examples and restrictions. Args: rt_input: The `RaggedTensor` from which a piece should be returned. key_list: The list of keys speci
(rt_input, key_list)
| 105 | |
| 106 | |
| 107 | def _ragged_getitem(rt_input, key_list): |
| 108 | """Helper for indexing and slicing ragged tensors with __getitem__(). |
| 109 | |
| 110 | Extracts the specified piece of the `rt_input`. See |
| 111 | `RaggedTensor.__getitem__` for examples and restrictions. |
| 112 | |
| 113 | Args: |
| 114 | rt_input: The `RaggedTensor` from which a piece should be returned. |
| 115 | key_list: The list of keys specifying which piece to return. Each key |
| 116 | corresponds with a separate dimension. |
| 117 | |
| 118 | Returns: |
| 119 | The indicated piece of rt_input. |
| 120 | |
| 121 | Raises: |
| 122 | ValueError: If `key_list` is not supported. |
| 123 | TypeError: If any keys in `key_list` have an unsupported type. |
| 124 | """ |
| 125 | if not key_list: |
| 126 | return rt_input |
| 127 | row_key = key_list[0] |
| 128 | inner_keys = key_list[1:] |
| 129 | |
| 130 | if row_key is Ellipsis: |
| 131 | expanded_key_list = _expand_ellipsis(key_list, rt_input.shape.ndims) |
| 132 | return _ragged_getitem(rt_input, expanded_key_list) |
| 133 | |
| 134 | # Adding a new axis: Get rt_input[inner_keys], and wrap it in a RaggedTensor |
| 135 | # that puts all values in a single row. |
| 136 | if row_key is array_ops.newaxis: |
| 137 | inner_rt = _ragged_getitem(rt_input, inner_keys) |
| 138 | nsplits = array_ops.shape(inner_rt.row_splits, |
| 139 | out_type=inner_rt.row_splits.dtype)[0] |
| 140 | return ragged_tensor.RaggedTensor.from_row_splits( |
| 141 | inner_rt, array_ops.stack([0, nsplits - 1]), validate=False) |
| 142 | |
| 143 | # Slicing a range of rows: first slice the outer dimension, and then |
| 144 | # call `_ragged_getitem_inner_dimensions` to handle the inner keys. |
| 145 | if isinstance(row_key, slice): |
| 146 | sliced_rt_input = _slice_ragged_row_dimension(rt_input, row_key) |
| 147 | return _ragged_getitem_inner_dimensions(sliced_rt_input, inner_keys) |
| 148 | |
| 149 | # Indexing a single row: slice values to get the indicated row, and then |
| 150 | # use a recursive call to __getitem__ to handle the inner keys. |
| 151 | else: |
| 152 | starts = rt_input.row_splits[:-1] |
| 153 | limits = rt_input.row_splits[1:] |
| 154 | if context.executing_eagerly(): |
| 155 | # In python, __getitem__ should throw IndexError for out of bound |
| 156 | # indices. This will allow iteration run correctly as python will |
| 157 | # translate IndexError into StopIteration for next()/__next__(). |
| 158 | # Below is an example: |
| 159 | # import tensorflow as tf |
| 160 | # r = tf.ragged.constant([[1., 2.], [3., 4., 5.], [6.]]) |
| 161 | # for elem in r: |
| 162 | # print(elem) |
| 163 | # In non eager mode, the exception is thrown when session runs |
| 164 | # so we don't know if out of bound happens before. |
no test coverage detected