def _gather(params, indices, axis, batch_dims):
"""Helper that implements the body for ragged gather().
Assumes that `params` and `indices` have been converted to tensors or
ragged tensors, and that `axis` and `batch_dims` have been normalized to
be positive. (So these conversions & normalizations can be skipped in
recursive calls to _gather).
Args:
params: The tensor from which to gather values.
indices: The indices of values to gather.
axis: The axis in `params` to gather `indices` from.
batch_dims: The number of batch dimensions.
Returns:
A potentially ragged tensor.
"""
params_is_ragged = ragged_tensor.is_ragged(params)
indices_is_ragged = ragged_tensor.is_ragged(indices)
if not (params_is_ragged or indices_is_ragged):
return array_ops.gather(params, indices, axis=axis, batch_dims=batch_dims)
if batch_dims > 0:
return _batch_gather(params, indices, axis, batch_dims)
if axis > 0:
return _axis_gather(params, indices, axis)
if indices_is_ragged:
return indices.with_values(_gather(params, indices.values, 0, 0))
if indices.shape.ndims is None:
raise ValueError('rank(indices) must be known statically')
out_ragged_rank = indices.shape.ndims + len(params.nested_row_splits) - 1
result = gen_ragged_array_ops.ragged_gather(
indices=indices,
params_dense_values=params.flat_values,
params_nested_splits=params.nested_row_splits,
OUTPUT_RAGGED_RANK=out_ragged_rank)
result = ragged_tensor.RaggedTensor.from_nested_row_splits(
result.output_dense_values, result.output_nested_splits, validate=False)
# Inject uniform_row_lengths into the result RaggedTensors for dimensions
# corresponding to dense outer dimensions of `indices`.
# TODO(edloper): Change this to construct the result using RowPartition
# objects instead, so we don't need to modify private variables.
if indices.shape.ndims > 1:
target = result
indices_shape = array_ops.shape(indices, out_type=params.row_splits.dtype)
shape_cumprod = math_ops.cumprod(indices_shape)
for dim in range(indices.shape.ndims - 1):
# pylint: disable=protected-access
target._cached_nrows = shape_cumprod[dim]
target._uniform_row_length = indices_shape[dim + 1]
target = target.values
return result
def gather(params, indices, validate_indices=None, axis=0, batch_dims=0,
name=None):
"""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` may be ragged. `indices` may be ragged.
`indices` must have dtype `int32` or `int64`. If any index is out of bounds,
then an error is returned.
Examples:
```python
>>> params = tf.constant(['a', 'b', 'c', 'd', 'e'])
>>> indices = tf.constant([3, 1, 2, 1, 0])
>>> ragged_params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
>>> ragged_indices = tf.ragged.constant([[3, 1, 2], [1], [], [0]])
>>> print ragged.gather(params, ragged_indices)
[['d', 'b', 'c'], ['b'], [], ['a']]
>>> print ragged.gather(ragged_params, indices)
[['e'], ['d'], [], ['d'], ['a', 'b', 'c']]
>>> print ragged.gather(ragged_params, ragged_indices)
[[['e'], ['d'], []], [['d']], [], [['a', 'b', 'c']]]
```
Args:
params: The potentially ragged tensor from which to gather values. Must be
at least rank 1.
indices: The potentially ragged tensor indicating which values to gather.
Must have dtype `int32` or `int64`. Values must be in the range `[0,
params.shape[0]]`.
validate_indices: Ignored.
axis: Must be zero.
batch_dims: Must be zero.
name: A name for the operation (optional).
Returns:
A `RaggedTensor`, where `output.dtype=params.dtype` and
`output.shape=indices.shape + params.shape[1:]` and
`output.ragged_rank=indices.shape.ndims + params.ragged_rank`.
Raises:
ValueError: If indices.shape.ndims is not known statically.
"""
del validate_indices
if not isinstance(axis, int) or axis != 0:
raise ValueError('axis != 0 is not supported for ragged gather yet.')
if not isinstance(batch_dims, int) or batch_dims != 0:
raise ValueError('batch_dims != 0 is not supported for ragged gather yet.')
with ops.name_scope(name, 'RaggedGather', [params, indices]):
params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
params, indices = ragged_tensor.match_row_splits_dtypes(params, indices)
if ragged_tensor.is_ragged(indices):
return indices.with_values(gather(params, indices.values))
if not ragged_tensor.is_ragged(params):
return array_ops.gather(params, indices)
indices = ops.convert_to_tensor(indices)
if indices.shape.ndims is None:
raise ValueError('indices.shape.ndims must be known statically')
result = gen_ragged_array_ops.ragged_gather(
indices=indices,
params_dense_values=params.flat_values,
params_nested_splits=params.nested_row_splits,
OUTPUT_RAGGED_RANK=indices.shape.ndims + len(params.nested_row_splits) -
1)
# Compose the RaggedTensor from splits & values.
return ragged_tensor.RaggedTensor.from_nested_row_splits(
result.output_dense_values, result.output_nested_splits, validate=False)
def gather(params, indices, validate_indices=None, axis=0, name=None):
"""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` may be ragged. `indices` may be ragged.
`indices` must have dtype `int32` or `int64`. If any index is out of bounds,
then an error is returned.
Examples:
```python
>>> params = tf.constant(['a', 'b', 'c', 'd', 'e'])
>>> indices = tf.constant([3, 1, 2, 1, 0])
>>> ragged_params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
>>> ragged_indices = tf.ragged.constant([[3, 1, 2], [1], [], [0]])
>>> print ragged.gather(params, ragged_indices)
[['d', 'b', 'c'], ['b'], [], ['a']]
>>> print ragged.gather(ragged_params, indices)
[['e'], ['d'], [], ['d'], ['a', 'b', 'c']]
>>> print ragged.gather(ragged_params, ragged_indices)
[[['e'], ['d'], []], [['d']], [], [['a', 'b', 'c']]]
```
Args:
params: The potentially ragged tensor from which to gather values. Must be
at least rank 1.
indices: The potentially ragged tensor indicating which values to gather.
Must have dtype `int32` or `int64`. Values must be in the range `[0,
params.shape[0]]`.
validate_indices: Ignored.
axis: Must be zero.
name: A name for the operation (optional).
Returns:
A `RaggedTensor`, where `output.dtype=params.dtype` and
`output.shape=indices.shape + params.shape[1:]` and
`output.ragged_rank=indices.shape.ndims + params.ragged_rank`.
Raises:
ValueError: If indices.shape.ndims is not known statically.
"""
del validate_indices
if not isinstance(axis, int) or axis != 0:
raise ValueError('axis>0 is not supported for ragged gather yet.')
with ops.name_scope(name, 'RaggedGather', [params, indices]):
params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
if ragged_tensor.is_ragged(indices):
return indices.with_values(gather(params, indices.values))
if not ragged_tensor.is_ragged(params):
return array_ops.gather(params, indices)
indices = ops.convert_to_tensor(indices)
if indices.shape.ndims is None:
raise ValueError('indices.shape.ndims must be known statically')
result = gen_ragged_array_ops.ragged_gather(
indices=indices,
params_dense_values=params.flat_values,
params_nested_splits=params.nested_row_splits,
OUTPUT_RAGGED_RANK=indices.shape.ndims + len(params.nested_row_splits) -
1)
# Compose the RaggedTensor from splits & values.
return ragged_tensor.RaggedTensor.from_nested_row_splits(
result.output_dense_values, result.output_nested_splits)
