def testGradientSegmentMax(self):
num_cols = 2
indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
num_segments = max(indices_flat) + 3
for indices in indices_flat, indices_flat.reshape(5, 2):
shape = indices.shape + (num_cols,)
with self.test_session(use_gpu=True):
tf_x, np_x = self._input(shape, dtype=dtypes_lib.float64)
s = math_ops.unsorted_segment_max(
data=tf_x, segment_ids=indices, num_segments=num_segments)
jacob_t, jacob_n = gradient_checker.compute_gradient(
tf_x,
shape,
s,
[num_segments, num_cols],
x_init_value=np_x.astype(np.double), delta=1)
self.assertAllClose(jacob_t, jacob_n)
def testGradientSegmentMax(self):
num_cols = 2
indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
num_segments = max(indices_flat) + 3
for indices in indices_flat, indices_flat.reshape(5, 2):
shape = indices.shape + (num_cols,)
with self.test_session(use_gpu=True):
tf_x, np_x = self._input(shape, dtype=dtypes_lib.float64)
s = math_ops.unsorted_segment_max(
data=tf_x, segment_ids=indices, num_segments=num_segments)
jacob_t, jacob_n = gradient_checker.compute_gradient(
tf_x,
shape,
s,
[num_segments, num_cols],
x_init_value=np_x.astype(np.double), delta=1)
self.assertAllClose(jacob_t, jacob_n)
def sparse_row_envelope(sparse_input, row_axis=0, col_axis=1, name=None):
"""Returns the length of each 'row' in a `SparseTensor`.
For example, if `sparse_input` has indices `[[0,0], [2, 0], [2, 1], [2, 2]]`
and shape `[3, 3]`, this function will return `[1, 0, 3]`.
Args:
sparse_input: a `SparseTensor` of rank at least 2.
row_axis: An integer. The axis for the row of the envelope matrix. Default
is 0.
col_axis: An integer. The axis for the col of the envelope matrix. Default
is 1.
name: A name for the operation (optional).
Returns:
A one-dimensional `Tensor` whose entries correspond to the length of each
row of `SparseTensor`.
Raises:
ValueError: If row_axis and col_axis are the same axis or they are not
integers.
"""
if not (isinstance(row_axis, compat.integral_types)
and isinstance(col_axis, compat.integral_types)):
raise ValueError("`row_axis` and `col_axis` must be integers.")
if row_axis == col_axis:
raise ValueError("Row and column can not be the same axis.")
with ops.name_scope(name, "sparse_row_envelope", [sparse_input]):
indices = sparse_input.indices
row_indices = indices[:, row_axis]
col_indices = indices[:, col_axis]
num_rows = math_ops.cast(sparse_input.dense_shape[row_axis],
dtypes.int32)
row_envelope = math_ops.unsorted_segment_max(col_indices + 1,
row_indices,
num_rows,
name=name)
zeros = array_ops.zeros_like(row_envelope)
return array_ops.where(row_envelope > zeros, row_envelope, zeros)
def sparse_row_envelope(sparse_input, row_axis=0, col_axis=1, name=None):
"""Returns the length of each 'row' in a `SparseTensor`.
For example, if `sparse_input` has indices `[[0,0], [2, 0], [2, 1], [2, 2]]`
and shape `[3, 3]`, this function will return `[1, 0, 3]`.
Args:
sparse_input: a `SparseTensor` of rank at least 2.
row_axis: An integer. The axis for the row of the envelope matrix. Default
is 0.
col_axis: An integer. The axis for the col of the envelope matrix. Default
is 1.
name: A name for the operation (optional).
Returns:
A one-dimensional `Tensor` whose entries correspond to the length of each
row of `SparseTensor`.
Raises:
ValueError: If row_axis and col_axis are the same axis or they are not
integers.
"""
if not (isinstance(row_axis, compat.integral_types) and
isinstance(col_axis, compat.integral_types)):
raise ValueError("`row_axis` and `col_axis` must be integers.")
if row_axis == col_axis:
raise ValueError("Row and column can not be the same axis.")
with ops.name_scope(name, "sparse_row_envelope", [sparse_input]):
indices = sparse_input.indices
row_indices = indices[:, row_axis]
col_indices = indices[:, col_axis]
num_rows = math_ops.cast(sparse_input.dense_shape[row_axis], dtypes.int32)
row_envelope = math_ops.unsorted_segment_max(
col_indices + 1, row_indices, num_rows, name=name)
zeros = array_ops.zeros_like(row_envelope)
return array_ops.where(row_envelope > zeros, row_envelope, zeros)
def _ragged_segment_aggregate(unsorted_segment_op, data, segment_ids,
num_segments, name=None):
"""Aggregates along segments of a RaggedTensor using `unsorted_segment_op`.
Returns a RaggedTensor `output` with `num_segments` rows, where the row
`output[i]` is formed by combining all rows of `data` whose corresponding
`segment_id` is `i`. The values in each row are combined using
`unsorted_segment_op`.
The length of the row `output[i]` will be the maximum of the lengths of
all rows of `data` whose corresponding `segment_id` is `i`. If no `data`
rows correspond to a given segment ID, then the output row for that segment
ID will be empty.
Args:
unsorted_segment_op: The tensorflow `op` that should be used to combine
values in each row. Must have the same signature and basic behavior as
`unsorted_segment_sum`, `unsorted_segment_max`, etc.
data: A `RaggedTensor` containing the values to be combined.
segment_ids: A `Tensor` or `RaggedTensor`. Must have type `int64` or
`int32`. `segment_ids.shape` must be a prefix of `data.shape`.
`segment_ids` is not required to be sorted.
num_segments: An `int32` or `int64` scalar.
name: A name prefix for the returned tensor (optional).
Returns:
A `RaggedTensor` containing the aggregated values. The returned tensor
has the same dtype as `data`, and its shape is
`[num_segments] + data.shape[segment_ids.rank:]`.
Raises:
ValueError: If segment_ids.shape is not a prefix of data.shape.
"""
if not (ragged_tensor.is_ragged(data) or
ragged_tensor.is_ragged(segment_ids)):
return unsorted_segment_op(data, segment_ids, num_segments, name)
with ops.name_scope(name, 'RaggedSegment',
[data, segment_ids, num_segments]) as name:
data = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
data, name='data')
segment_ids = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
segment_ids, name='segment_ids')
if ragged_tensor.is_ragged(segment_ids):
if not ragged_tensor.is_ragged(data):
raise ValueError('segment_ids.shape must be a prefix of data.shape, '
'but segment_ids is ragged and data is not.')
check_splits = check_ops.assert_equal(
segment_ids.row_splits,
data.row_splits,
message='segment_ids.shape must be a prefix of data.shape')
with ops.control_dependencies([check_splits]):
return _ragged_segment_aggregate(unsorted_segment_op, data.values,
segment_ids.values, num_segments, name)
segment_ids = math_ops.cast(segment_ids, dtypes.int64)
# Find the length of each row in data. (dtype=int64, shape=[data_nrows])
data_row_lengths = data.row_splits[1:] - data.row_splits[:-1]
# Find the length that each output row will have. The length of the row
# corresponding to segment `id` is `max(data_row_lengths[i])` where
# `segment_ids[i]=id`. (dtype=int64, shape=[output_nrows])
output_row_lengths = math_ops.maximum(
math_ops.unsorted_segment_max(data_row_lengths, segment_ids,
num_segments), 0)
assert output_row_lengths.dtype == dtypes.int64
# Build the splits tensor for the output RaggedTensor.
output_splits = array_ops.concat(
[
array_ops.zeros([1], dtypes.int64),
math_ops.cumsum(output_row_lengths)
],
axis=0)
# For each row in `data`, find the start & limit position where that row's
# values will be aggregated in output.values.
data_row_to_out_row_start = array_ops.gather(output_splits, segment_ids)
data_row_to_out_row_limit = data_row_to_out_row_start + data_row_lengths
# For each value in `data.values`, find the position where it will
# aggregated in `output.values`.
# Get the target output values index for each data values index.
data_val_to_out_val_index = range(data_row_to_out_row_start,
data_row_to_out_row_limit).values
# Recursively aggregate the values.
output_values = _ragged_segment_aggregate(unsorted_segment_op, data.values,
data_val_to_out_val_index,
output_splits[-1])
return ragged_factory_ops.from_row_splits(output_values, output_splits)
def _ragged_segment_aggregate(unsorted_segment_op,
data,
segment_ids,
num_segments,
separator=None,
name=None):
"""Aggregates along segments of a RaggedTensor using `unsorted_segment_op`.
Returns a RaggedTensor `output` with `num_segments` rows, where the row
`output[i]` is formed by combining all rows of `data` whose corresponding
`segment_id` is `i`. The values in each row are combined using
`unsorted_segment_op`.
The length of the row `output[i]` will be the maximum of the lengths of
all rows of `data` whose corresponding `segment_id` is `i`. If no `data`
rows correspond to a given segment ID, then the output row for that segment
ID will be empty.
Args:
unsorted_segment_op: The tensorflow `op` that should be used to combine
values in each row. Must have the same signature and basic behavior as
`unsorted_segment_sum`, `unsorted_segment_max`, etc.
data: A `RaggedTensor` containing the values to be combined.
segment_ids: A `Tensor` or `RaggedTensor`. Must have type `int64` or
`int32`. `segment_ids.shape` must be a prefix of `data.shape`.
`segment_ids` is not required to be sorted.
num_segments: An `int32` or `int64` scalar.
separator: An optional string. Defaults to None. The separator to use when
joining. Only used for string types.
name: A name prefix for the returned tensor (optional).
Returns:
A `RaggedTensor` containing the aggregated values. The returned tensor
has the same dtype as `data`, and its shape is
`[num_segments] + data.shape[segment_ids.rank:]`.
Raises:
ValueError: If segment_ids.shape is not a prefix of data.shape.
"""
if not (ragged_tensor.is_ragged(data) or
ragged_tensor.is_ragged(segment_ids)):
if separator is not None:
# It uses unsorted_segment_join.
return unsorted_segment_op(data, segment_ids, num_segments, separator,
name)
else:
return unsorted_segment_op(data, segment_ids, num_segments, name)
with ops.name_scope(name, 'RaggedSegment',
[data, segment_ids, num_segments]) as name:
data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
segment_ids = ragged_tensor.convert_to_tensor_or_ragged_tensor(
segment_ids, name='segment_ids')
data, segment_ids = ragged_tensor.match_row_splits_dtypes(data, segment_ids)
if segment_ids.dtype not in (dtypes.int32, dtypes.int64):
raise ValueError('segment_ids must have dtype int32 or int64.')
if ragged_tensor.is_ragged(segment_ids):
if not ragged_tensor.is_ragged(data):
raise ValueError('segment_ids.shape must be a prefix of data.shape, '
'but segment_ids is ragged and data is not.')
check_splits = check_ops.assert_equal(
segment_ids.row_splits,
data.row_splits,
message='segment_ids.shape must be a prefix of data.shape')
with ops.control_dependencies([check_splits]):
return _ragged_segment_aggregate(unsorted_segment_op, data.values,
segment_ids.values, num_segments,
separator)
# Find the length of each row in data. (shape=[data_nrows])
data_row_lengths = data.row_splits[1:] - data.row_splits[:-1]
# Find the length that each output row will have. The length of the row
# corresponding to segment `id` is `max(data_row_lengths[i])` where
# `segment_ids[i]=id`. (shape=[output_nrows])
output_row_lengths = math_ops.maximum(
math_ops.unsorted_segment_max(data_row_lengths, segment_ids,
num_segments), 0)
# Build the splits tensor for the output RaggedTensor.
output_splits = array_ops.concat([
array_ops.zeros([1], output_row_lengths.dtype),
math_ops.cumsum(output_row_lengths)
],
axis=0)
# For each row in `data`, find the start & limit position where that row's
# values will be aggregated in output.values.
data_row_to_out_row_start = array_ops.gather(output_splits, segment_ids)
data_row_to_out_row_limit = data_row_to_out_row_start + data_row_lengths
# For each value in `data.values`, find the position where it will
# aggregated in `output.values`.
# Get the target output values index for each data values index.
data_val_to_out_val_index = range(data_row_to_out_row_start,
data_row_to_out_row_limit).values
# Recursively aggregate the values.
output_values = _ragged_segment_aggregate(unsorted_segment_op, data.values,
data_val_to_out_val_index,
output_splits[-1], separator)
return ragged_tensor.RaggedTensor.from_row_splits(
output_values, output_splits, validate=False)
