Ragged version of tf.where(condition, x, y).
(condition, x, y)
| 112 | |
| 113 | |
| 114 | def _elementwise_where(condition, x, y): |
| 115 | """Ragged version of tf.where(condition, x, y).""" |
| 116 | condition_is_ragged = isinstance(condition, ragged_tensor.RaggedTensor) |
| 117 | x_is_ragged = isinstance(x, ragged_tensor.RaggedTensor) |
| 118 | y_is_ragged = isinstance(y, ragged_tensor.RaggedTensor) |
| 119 | |
| 120 | if not (condition_is_ragged or x_is_ragged or y_is_ragged): |
| 121 | return array_ops.where(condition, x, y) |
| 122 | |
| 123 | elif condition_is_ragged and x_is_ragged and y_is_ragged: |
| 124 | return ragged_functional_ops.map_flat_values(array_ops.where, condition, x, |
| 125 | y) |
| 126 | elif not condition_is_ragged: |
| 127 | # Concatenate x and y, and then use `gather` to assemble the selected rows. |
| 128 | condition.shape.assert_has_rank(1) |
| 129 | x_and_y = ragged_concat_ops.concat([x, y], axis=0) |
| 130 | x_nrows = _nrows(x, out_type=x_and_y.row_splits.dtype) |
| 131 | y_nrows = _nrows(y, out_type=x_and_y.row_splits.dtype) |
| 132 | indices = array_ops.where(condition, math_ops.range(x_nrows), |
| 133 | x_nrows + math_ops.range(y_nrows)) |
| 134 | return ragged_gather_ops.gather(x_and_y, indices) |
| 135 | |
| 136 | else: |
| 137 | raise ValueError('Input shapes do not match.') |
| 138 | |
| 139 | |
| 140 | def _coordinate_where(condition): |