def testShape(self):
self.assertAllEqual(
ragged_math_ops.range(0, 0, 1).shape.as_list(), [1, None])
self.assertAllEqual(
ragged_math_ops.range([1, 2, 3]).shape.as_list(), [3, None])
self.assertAllEqual(
ragged_math_ops.range([1, 2, 3], [4, 5, 6]).shape.as_list(), [3, None])
def testKernelErrors(self):
with self.assertRaisesRegex(errors.InvalidArgumentError,
r'Requires delta != 0'):
self.evaluate(ragged_math_ops.range(0, 0, 0))
with self.assertRaisesRegex(errors.InvalidArgumentError,
r'Requires \(\(limit - start\) / delta\) <='):
self.evaluate(ragged_math_ops.range(0.1, 1e10, 1e-10))
def testRaggedWithDifferentShapes(self):
dataset = dataset_ops.Dataset.range(10).map(ragged_math_ops.range).batch(5)
expected_output = [
ragged_concat_ops.stack([ragged_math_ops.range(i) for i in range(5)]),
ragged_concat_ops.stack(
[ragged_math_ops.range(i) for i in range(5, 10)])
]
self.assertDatasetProduces(dataset, expected_output=expected_output)
def testBatchRaggedWithDifferentShapes(self):
dataset = dataset_ops.Dataset.range(10).map(ragged_math_ops.range).batch(5)
expected_output = [
ragged_concat_ops.stack([ragged_math_ops.range(i) for i in range(5)]),
ragged_concat_ops.stack(
[ragged_math_ops.range(i) for i in range(5, 10)])
]
self.assertDatasetProduces(dataset, expected_output=expected_output)
def testDocStringExamples(self):
"""Examples from ragged_range.__doc__."""
rt1 = ragged_math_ops.range([3, 5, 2])
self.assertAllEqual(rt1, [[0, 1, 2], [0, 1, 2, 3, 4], [0, 1]])
rt2 = ragged_math_ops.range([0, 5, 8], [3, 3, 12])
self.assertAllEqual(rt2, [[0, 1, 2], [], [8, 9, 10, 11]])
rt3 = ragged_math_ops.range([0, 5, 8], [3, 3, 12], 2)
self.assertAllEqual(rt3, [[0, 2], [], [8, 10]])
def testBroadcast(self):
# Specify starts and limits, broadcast deltas.
self.assertAllEqual(
ragged_math_ops.range([0, 3, 5], [4, 4, 15], 3),
[list(range(0, 4, 3)),
list(range(3, 4, 3)),
list(range(5, 15, 3))])
# Broadcast all arguments.
self.assertAllEqual(ragged_math_ops.range(0, 5, 1), [list(range(0, 5, 1))])
def testNegativeDeltas(self):
self.assertAllEqual(
ragged_math_ops.range([0, 3, 5], limits=0, deltas=-1),
[list(range(0, 0, -1)),
list(range(3, 0, -1)),
list(range(5, 0, -1))])
self.assertAllEqual(
ragged_math_ops.range([0, -3, 5], limits=0, deltas=[-1, 1, -2]),
[list(range(0, 0, -1)),
list(range(-3, 0, 1)),
list(range(5, 0, -2))])
def testFloatRanges(self):
expected = [[0.0, 0.4, 0.8, 1.2, 1.6, 2.0, 2.4, 2.8, 3.2, 3.6], [3.0],
[5.0, 7.2, 9.4, 11.6, 13.8]]
actual = ragged_math_ops.range([0.0, 3.0, 5.0], [3.9, 4.0, 15.0],
[0.4, 1.5, 2.2])
self.assertEqual(expected, [[round(v, 5) for v in row]
for row in self.eval_to_list(actual)])
def _ragged_gather_grad(op, *grads):
"""Gradient for RaggedGather op."""
param_nested_splits = op.inputs[:-2]
param_inner_values = op.inputs[-2]
indices = op.inputs[-1]
grad_inner_values = grads[-1]
# For each row in `params`, find the range of values in `params.inner_values`
# that is covered by that row. In particular, the values in row `i` are
# `param_inner_values[combined_splits[i]:combined_splits[i+1]`.
combined_splits = param_nested_splits[0]
for row_splits in param_nested_splits[1:]:
combined_splits = array_ops.gather(row_splits, combined_splits)
# The outer dimensions of `indices` correspond 1:1 with the outer dimensions
# of `ragged_grad` that are encoded by `grad_nested_splits`. Thus, the
# flattened `indices` correspond 1:1 with `grad_inner_values`.
flat_indices = array_ops.reshape(indices, [-1])
# Build an IndexedSlices where the values are taken from `flat_grad`.
grad_indices = ragged_math_ops.range(
array_ops.gather(combined_splits, flat_indices),
array_ops.gather(combined_splits[1:], flat_indices)).values
param_inner_values_grad = indexed_slices.IndexedSlices(
values=grad_inner_values,
indices=grad_indices,
dense_shape=array_ops.shape(param_inner_values))
return [None
for _ in param_nested_splits] + [param_inner_values_grad, None]
def _get_selected_item_positions(item_selector, input_ids, axis=1):
"""Get the positions of the items that have been selected.
Args:
item_selector: an instance of `ItemSelector`.
input_ids: a `RaggedTensor` with n dimensions, whose items will be
selected on.
axis: (optional) An int detailing the dimension to apply selection on.
Default is the 1st dimension.
Returns:
A `RaggedTensor` of int64s, with rank 2, shape
[batch, (num_selections)] and whose values are the positions of items
that have been selected.
"""
original_input_ids = input_ids
# select items for masking
selected_for_mask = item_selector.get_selection_mask(input_ids, axis)
# create a positions RT
original_input_ids = (original_input_ids.merge_dims(
1, -1) if original_input_ids.ragged_rank > 1 else original_input_ids)
positions = ragged_math_ops.range(original_input_ids.row_lengths())
positions = input_ids.with_flat_values(positions.flat_values)
# drop out not-masked positions
results = ragged_array_ops.boolean_mask(positions, selected_for_mask)
results = results.merge_dims(1, -1) if results.ragged_rank > 1 else results
return results
def _get_selection_mask(original, num_to_select, axis=-1):
"""Get a selection mask given how many items to select."""
num_to_select = ops.convert_to_tensor(num_to_select)
num_to_select = array_ops.reshape(num_to_select, [-1])
row_lengths = _get_row_lengths_merged_to_axis(original, axis)
num_to_select = array_ops.broadcast_to(num_to_select,
array_ops.shape(row_lengths))
num_to_select = math_ops.cast(num_to_select, row_lengths.dtype)
num_to_select = math_ops.minimum(num_to_select, row_lengths)
ones = array_ops.ones_like(ragged_math_ops.range(num_to_select))
ones = math_ops.cast(ones, dtypes.int32)
zeros_row_length = row_lengths - num_to_select
zeros = math_ops.cast(
array_ops.zeros_like(ragged_math_ops.range(zeros_row_length)),
dtypes.int32)
results = array_ops.concat([ones, zeros], 1)
results = math_ops.cast(results, dtypes.bool)
return results
def testBasicRanges(self):
# Specify limits only.
self.assertAllEqual(
ragged_math_ops.range([0, 3, 5]),
[list(range(0)), list(range(3)),
list(range(5))])
# Specify starts and limits.
self.assertAllEqual(
ragged_math_ops.range([0, 3, 5], [2, 3, 10]),
[list(range(0, 2)),
list(range(3, 3)),
list(range(5, 10))])
# Specify starts, limits, and deltas.
self.assertAllEqual(
ragged_math_ops.range([0, 3, 5], [4, 4, 15], [2, 3, 4]),
[list(range(0, 4, 2)),
list(range(3, 4, 3)),
list(range(5, 15, 4))])
def _build_ragged_tensor_from_value_ranges(starts, limits, step, values):
"""Returns a `RaggedTensor` containing the specified sequences of values.
Returns a RaggedTensor `output` where:
```python
output.shape[0] = starts.shape[0]
output[i] = values[starts[i]:limits[i]:step]
```
Requires that `starts.shape == limits.shape` and
`0 <= starts[i] <= limits[i] <= values.shape[0]`.
Args:
starts: 1D integer Tensor specifying the start indices for the sequences of
values to include.
limits: 1D integer Tensor specifying the limit indices for the sequences of
values to include.
step: Integer value specifying the step size for strided slices.
values: The set of values to select from.
Returns:
A `RaggedTensor`.
Raises:
ValueError: Until the prerequisite ops are checked in.
"""
# Use `ragged_range` to get the index of each value we should include.
if step is None:
step = 1
step = ops.convert_to_tensor(step, name="step")
if step.dtype.is_integer:
step = math_ops.cast(step, starts.dtype)
else:
raise TypeError("slice strides must be integers or None")
value_indices = ragged_math_ops.range(starts,
limits,
step,
row_splits_dtype=starts.dtype)
# Use `ragged_gather` or `array_ops.gather` to collect the values.
if isinstance(values, ragged_tensor.RaggedTensor):
gathered_values = ragged_gather_ops.gather(
params=values, indices=value_indices.values)
else:
gathered_values = array_ops.gather(params=values,
indices=value_indices.values)
# Assemble the RaggedTensor from splits & values.
return value_indices.with_values(gathered_values)
def _build_ragged_tensor_from_value_ranges(starts, limits, step, values):
"""Returns a `RaggedTensor` containing the specified sequences of values.
Returns a RaggedTensor `output` where:
```python
output.shape[0] = starts.shape[0]
output[i] = values[starts[i]:limits[i]:step]
```
Requires that `starts.shape == limits.shape` and
`0 <= starts[i] <= limits[i] <= values.shape[0]`.
Args:
starts: 1D integer Tensor specifying the start indices for the sequences of
values to include.
limits: 1D integer Tensor specifying the limit indices for the sequences of
values to include.
step: Integer value specifying the step size for strided slices.
values: The set of values to select from.
Returns:
A `RaggedTensor`.
Raises:
ValueError: Until the prerequisite ops are checked in.
"""
# Use `ragged_range` to get the index of each value we should include.
if step is None:
step = 1
step = ops.convert_to_tensor(step, name="step")
if step.dtype.is_integer:
step = math_ops.cast(step, dtypes.int64)
else:
raise TypeError("slice strides must be integers or None")
value_indices = ragged_math_ops.range(starts, limits, step)
# Use `ragged_gather` or `array_ops.gather` to collect the values.
if isinstance(values, ragged_tensor.RaggedTensor):
gathered_values = ragged_gather_ops.gather(
params=values, indices=value_indices.values)
else:
gathered_values = array_ops.gather(
params=values, indices=value_indices.values)
# Assemble the RaggedTensor from splits & values.
return value_indices.with_values(gathered_values)
def testEmptyRanges(self):
rt1 = ragged_math_ops.range([0, 5, 3], [0, 3, 5])
rt2 = ragged_math_ops.range([0, 5, 5], [0, 3, 5], -1)
self.assertAllEqual(rt1, [[], [], [3, 4]])
self.assertAllEqual(rt2, [[], [5, 4], []])
def testFloatRanges(self):
expected = [[0.0, 0.4, 0.8, 1.2, 1.6, 2.0, 2.4, 2.8, 3.2, 3.6], [3.0],
[5.0, 7.2, 9.4, 11.6, 13.8]]
actual = ragged_math_ops.range([0.0, 3.0, 5.0], [3.9, 4.0, 15.0],
[0.4, 1.5, 2.2])
self.assertAllClose(actual, expected)
def testKernelErrors(self):
with self.assertRaisesRegexp(errors.InvalidArgumentError,
r'Requires delta != 0'):
self.evaluate(ragged_math_ops.range(0, 0, 0))
class RaggedMapOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
@parameterized.parameters([
# The following test sets map over a RaggedTensor and apply a
# transformation that returns with shape:
# [d1, (d2)] -> [d1]
dict(
fn=mo.reduce_mean,
elems=[[1, 2, 3], [4, 5], [6, 7]],
elems_dtype=dtypes.int32,
expected_output=[2, 4, 6],
result_dtype=dtypes.int32,
),
dict(
fn=string_ops.reduce_join,
elems=[['foo', 'bar', 'baz'], ['a'], ['b', 'c']],
expected_output=[b'foobarbaz', b'a', b'bc'],
elems_dtype=dtypes.string,
result_dtype=dtypes.string,
),
# [d1, (d2)] -> [d1, 2]
dict(
fn=lambda x: array_ops.stack([mo.reduce_mean(x),
mo.reduce_sum(x)]),
# fn=self.stack_mean_and_sum,
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 6], [4.5, 9], [6.5, 13]],
elems_dtype=dtypes.float32,
result_dtype=dtypes.float32,
expected_ragged_rank=0,
),
# [d1, (d2)] -> [d1, (d2)]
dict(
fn=lambda x: x + np.int64(1),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 3, 4], [5, 6], [7, 8]],
elems_dtype=dtypes.int64,
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), d3] -> [d1, (d2), d3]
dict(
fn=lambda x: x + np.int64(1),
elems=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
elems_ragged_rank=1,
expected_ragged_rank=1,
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
expected_output=[[[2, 3], [4, 5]], [], [[6, 7], [8, 9], [10, 1]]],
),
# [d1, (d2)] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0]),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[[1, 2, 3]], [[4, 5]], [[6, 7]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_functional_ops.map_flat_values(mo.add, x, 1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[[2, 3, 4]], [[5, 6], [7, 8]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2)]
dict(
fn=lambda x: ragged_math_ops.reduce_sum(x, axis=1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[6], [9, 13]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1, (d3)]
dict(
fn=lambda x: ragged_math_ops.reduce_sum(x, axis=0),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[1, 2, 3], [10, 12]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1]
dict(
fn=ragged_math_ops.reduce_sum,
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[6, 22],
result_dtype=dtypes.int64,
),
# [d1] -> [d1, (d2)]
dict(
fn=mo.range,
elems=[4, 0, 2],
expected_output=[[0, 1, 2, 3], [], [0, 1]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_math_ops.range(mo.range(x)),
elems=[5, 0, 3],
expected_output=[[[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]], [],
[[], [0], [0, 1]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3), (d4a), (d5)] -> [d1, (d2), (d3), (d4b), (d5)]
dict(
fn=lambda x: x + np.int64(1),
elems=[[[[[1, 2, 3]], [[4], [5]]]], [[[[6, 7]]], [[[8], []]]]],
expected_output=[[[[[2, 3, 4]], [[5], [6]]]],
[[[[7, 8]]], [[[9], []]]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=4),
),
])
def testRaggedMap(
self,
fn,
elems,
expected_output,
expected_ragged_rank=None,
result_ragged_rank=None,
elems_ragged_rank=None,
elems_dtype=dtypes.int64,
result_dtype=None,
infer_shape=True,
):
elems = ragged_factory_ops.constant(elems, elems_dtype,
elems_ragged_rank)
output = ragged_map_ops.map_fn(fn=fn,
elems=elems,
dtype=result_dtype,
infer_shape=infer_shape)
expected_rt = ragged_factory_ops.constant(
expected_output, ragged_rank=expected_ragged_rank)
self.assertAllEqual(expected_rt, output)
def testRaggedMapOnStructure(self):
batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _reduce_sum_from_all(f):
return mo.reduce_sum(f['batman']) + mo.reduce_sum(f['robin'])
output = ragged_map_ops.map_fn(
fn=_reduce_sum_from_all,
elems=features,
dtype=dtypes.int32,
)
self.assertAllEqual(output, [66, 44, 198])
# Test mapping over a dict of RTs can produce a dict of RTs.
def testRaggedMapOnStructure_RaggedOutputs(self):
batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _increment(f):
return {
'batman': f['batman'] + 1,
'robin': f['robin'] + 1,
}
output = ragged_map_ops.map_fn(
fn=_increment,
elems=features,
infer_shape=False,
dtype={
'batman':
ragged_tensor.RaggedTensorType(dtype=dtypes.int32,
ragged_rank=1),
'robin':
ragged_tensor.RaggedTensorType(dtype=dtypes.int32,
ragged_rank=1)
},
)
self.assertAllEqual(output['batman'], [[2, 3, 4], [5], [6, 7, 8]])
self.assertAllEqual(output['robin'],
[[11, 21, 31], [41], [51, 61, 71]])
def testZip(self):
x = ragged_factory_ops.constant(
[[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64)
y = array_ops.expand_dims(mo.range(x.nrows(out_type=dtypes.int64)),
axis=1)
def _zip(foo):
y_val, x_val = foo
bar = array_ops.tile(y_val, array_ops.shape(x_val))
return array_ops.stack([bar, x_val], axis=1)
output = ragged_map_ops.map_fn(_zip, (y, x),
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=1),
infer_shape=False)
self.assertAllEqual(
output,
[[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]],
[[3, 70]], [[4, 80], [4, 90], [4, 100]]])
def testBatchGather(self):
tokens = ragged_factory_ops.constant([['hello', '.', 'there'],
['merhaba'],
['bonjour', '.', 'ca va', '?']])
indices = ragged_factory_ops.constant([[0, 2], [0], [0, 2]])
def gather(x):
tokens_val, indices_val = x
return array_ops.gather(tokens_val, indices_val)
data = tokens, indices
out = ragged_map_ops.map_fn(gather,
data,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.string, ragged_rank=1),
infer_shape=False)
self.assertAllEqual(
out, [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
def testMismatchRaggedRank(self):
elems = ragged_factory_ops.constant([[[1, 2, 3]], [[4, 5], [6, 7]]])
fn = lambda x: ragged_math_ops.reduce_sum(x, axis=0)
with self.assertRaisesRegex(
ValueError, r'(?s)Expected `fn` to return.*But it returned.*'):
_ = ragged_map_ops.map_fn(fn,
elems,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=23))
def testMismatchRaggedRank2(self):
elems = ragged_factory_ops.constant([[1, 2, 3], [4, 5], [6, 7]])
fn = lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0])
with self.assertRaisesRegex(
ValueError, r'(?s)Expected `fn` to return.*But it returned.*'):
_ = ragged_map_ops.map_fn(fn,
elems,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=10))
def testMapOnSparseTensor(self):
s = sparse_tensor.SparseTensor(
indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
values=[0, 5, 0, 4],
dense_shape=[2, 2],
)
t2 = ragged_tensor.RaggedTensor.from_sparse(s)
id_t2 = ragged_map_ops.map_fn(
lambda x: x,
t2,
)
self.assertAllEqual(id_t2, [[0, 5], [0, 4]])
