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)