Ragged version of tf.where(condition).
(condition)
| 138 | |
| 139 | |
| 140 | def _coordinate_where(condition): |
| 141 | """Ragged version of tf.where(condition).""" |
| 142 | if not isinstance(condition, ragged_tensor.RaggedTensor): |
| 143 | return array_ops.where(condition) |
| 144 | |
| 145 | # The coordinate for each `true` value in condition.values. |
| 146 | selected_coords = _coordinate_where(condition.values) |
| 147 | |
| 148 | # Convert the first index in each coordinate to a row index and column index. |
| 149 | condition = condition.with_row_splits_dtype(selected_coords.dtype) |
| 150 | first_index = selected_coords[:, 0] |
| 151 | selected_rows = array_ops.gather(condition.value_rowids(), first_index) |
| 152 | selected_row_starts = array_ops.gather(condition.row_splits, selected_rows) |
| 153 | selected_cols = first_index - selected_row_starts |
| 154 | |
| 155 | # Assemble the row & column index with the indices for inner dimensions. |
| 156 | return array_ops.concat([ |
| 157 | array_ops.expand_dims(selected_rows, 1), |
| 158 | array_ops.expand_dims(selected_cols, 1), selected_coords[:, 1:] |
| 159 | ], |
| 160 | axis=1) |
| 161 | |
| 162 | |
| 163 | def _nrows(rt_input, out_type): |
no test coverage detected