def ragged_one_hot(indices: ragged_tensor.Ragged,
depth,
on_value=None,
off_value=None,
axis=None,
dtype=None,
name=None):
"""Applies tf.one_hot along the values of a RaggedTensor."""
# Get the adjusted axis value for the call to array_ops.one_hot.
# Note: the only negative `axis` value supported by array_ops.one_hot is -1.
if isinstance(axis, int) and axis >= 0:
if axis <= indices.ragged_rank:
raise ValueError('axis (%d) must be greater than indices.ragged_rank '
'(%d).' % (axis, indices.ragged_rank))
axis -= indices.ragged_rank
with ops.name_scope(name, 'RaggedOneHot',
[indices, depth, on_value, off_value, axis]):
indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
return indices.with_flat_values(
array_ops.one_hot(indices.flat_values, depth, on_value, off_value, axis,
dtype, name))
def expand_dims(input: ragged_tensor.Ragged, axis, name=None): # pylint: disable=redefined-builtin
"""Inserts a dimension with shape 1 into a potentially ragged tensor's shape.
Given a potentially ragged tenor `input`, this operation inserts a
dimension with size 1 at the dimension `axis` of `input`'s shape.
The following table gives some examples showing how `ragged.expand_dims`
impacts the shapes of different input tensors. Ragged dimensions are
indicated by enclosing them in parentheses.
input.shape | axis | result.shape
----------------------- | ---- | -----------------------------
`[D1, D2]` | `0` | `[1, D1, D2]`
`[D1, D2]` | `1` | `[D1, 1, D2]`
`[D1, D2]` | `2` | `[D1, D2, 1]`
`[D1, (D2), (D3), D4]` | `0` | `[1, D1, (D2), (D3), D4]`
`[D1, (D2), (D3), D4]` | `1` | `[D1, 1, (D2), (D3), D4]`
`[D1, (D2), (D3), D4]` | `2` | `[D1, (D2), 1, (D3), D4]`
`[D1, (D2), (D3), D4]` | `3` | `[D1, (D2), (D3), 1, D4]`
`[D1, (D2), (D3), D4]` | `4` | `[D1, (D2), (D3), D4, 1]`
Args:
input: The potentially tensor that should be expanded with a new dimension.
axis: An integer constant indicating where the new dimension should be
inserted.
name: A name for the operation (optional).
Returns:
A tensor with the same values as `input`, with an added dimension of
size 1 at `axis`.
#### Examples:
>>> rt = tf.ragged.constant([[1, 2], [3]])
>>> print(rt.shape)
(2, None)
>>> expanded = tf.expand_dims(rt, axis=0)
>>> print(expanded.shape, expanded)
(1, 2, None) <tf.RaggedTensor [[[1, 2], [3]]]>
>>> expanded = tf.expand_dims(rt, axis=1)
>>> print(expanded.shape, expanded)
(2, 1, None) <tf.RaggedTensor [[[1, 2]], [[3]]]>
>>> expanded = tf.expand_dims(rt, axis=2)
>>> print(expanded.shape, expanded)
(2, None, 1) <tf.RaggedTensor [[[1], [2]], [[3]]]>
"""
with ops.name_scope(name, 'RaggedExpandDims', [input]):
input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
input, name='input')
if not ragged_tensor.is_ragged(input):
return array_ops.expand_dims(input, axis)
ndims = None if input.shape.ndims is None else input.shape.ndims + 1
axis = array_ops.get_positive_axis(axis, ndims, ndims_name='rank(input)')
if axis == 0:
return ragged_tensor.RaggedTensor.from_uniform_row_length(
input, uniform_row_length=input.nrows(), nrows=1, validate=False)
elif axis == 1:
return ragged_tensor.RaggedTensor.from_uniform_row_length(
input, uniform_row_length=1, nrows=input.nrows(), validate=False)
else:
if ragged_tensor.is_ragged(input.values):
return input.with_values(expand_dims(input.values, axis - 1))
else:
return input.with_values(array_ops.expand_dims(input.values, axis - 1))
def squeeze(input: ragged_tensor.Ragged, axis=None, name=None): # pylint: disable=redefined-builtin
"""Ragged compatible squeeze.
If `input` is a `tf.Tensor`, then this calls `tf.squeeze`.
If `input` is a `tf.RaggedTensor`, then this operation takes `O(N)` time,
where `N` is the number of elements in the squeezed dimensions.
Args:
input: A potentially ragged tensor. The input to squeeze.
axis: An optional list of ints. Defaults to `None`. If the `input` is
ragged, it only squeezes the dimensions listed. It fails if `input` is
ragged and axis is []. If `input` is not ragged it calls tf.squeeze. Note
that it is an error to squeeze a dimension that is not 1. It must be in
the range of [-rank(input), rank(input)).
name: A name for the operation (optional).
Returns:
A potentially ragged tensor. Contains the same data as input,
but has one or more dimensions of size 1 removed.
"""
with ops.name_scope(name, 'RaggedSqueeze', [input]):
input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input)
if isinstance(input, ops.Tensor):
return array_ops.squeeze(input, axis, name)
if axis is None:
raise ValueError('Ragged.squeeze must have an axis argument.')
if isinstance(axis, int):
axis = [axis]
elif ((not isinstance(axis, (list, tuple)))
or (not all(isinstance(d, int) for d in axis))):
raise TypeError('Axis must be a list or tuple of integers.')
dense_dims = []
ragged_dims = []
# Normalize all the dims in axis to be positive
axis = [
array_ops.get_positive_axis(d, input.shape.ndims, 'axis[%d]' % i,
'rank(input)')
for i, d in enumerate(axis)
]
for dim in axis:
if dim > input.ragged_rank:
dense_dims.append(dim - input.ragged_rank)
else:
ragged_dims.append(dim)
# Make sure the specified ragged dimensions are squeezable.
assertion_list = []
scalar_tensor_one = constant_op.constant(1,
dtype=input.row_splits.dtype)
for i, r in enumerate(input.nested_row_lengths()):
if i + 1 in ragged_dims:
assertion_list.append(
control_flow_ops.Assert(
math_ops.reduce_all(
math_ops.equal(r, scalar_tensor_one)),
[
'the given axis (axis = %d) is not squeezable!' %
(i + 1)
]))
if 0 in ragged_dims:
scalar_tensor_two = constant_op.constant(2, dtype=dtypes.int32)
assertion_list.append(
control_flow_ops.Assert(
math_ops.equal(array_ops.size(input.row_splits),
scalar_tensor_two),
['the given axis (axis = 0) is not squeezable!']))
# Till now, we are sure that the ragged dimensions are squeezable.
squeezed_rt = None
squeezed_rt = control_flow_ops.with_dependencies(
assertion_list, input.flat_values)
if dense_dims:
# Gives error if the dense dimension is not squeezable.
squeezed_rt = array_ops.squeeze(squeezed_rt, dense_dims)
remaining_row_splits = []
remaining_row_splits = list()
for i, row_split in enumerate(input.nested_row_splits):
# each row_splits tensor is for dimension #(i+1) .
if (i + 1) not in ragged_dims:
remaining_row_splits.append(row_split)
# Take care of the first row if it is to be squeezed.
if remaining_row_splits and 0 in ragged_dims:
remaining_row_splits.pop(0)
squeezed_rt = RaggedTensor.from_nested_row_splits(
squeezed_rt, remaining_row_splits)
# Corner case: when removing all the ragged dimensions and the output is
# a scalar tensor e.g. ragged.squeeze(ragged.constant([[[1]]])).
if set(range(0, input.ragged_rank + 1)).issubset(set(ragged_dims)):
squeezed_rt = array_ops.squeeze(squeezed_rt, [0], name)
return squeezed_rt