1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Ops for converting between row_splits and segment_ids.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21from tensorflow.python.framework import dtypes 22from tensorflow.python.framework import ops 23from tensorflow.python.framework import tensor_shape 24from tensorflow.python.framework import tensor_util 25from tensorflow.python.ops import array_ops 26from tensorflow.python.ops import math_ops 27from tensorflow.python.ops.ragged import ragged_util 28from tensorflow.python.util import dispatch 29from tensorflow.python.util.tf_export import tf_export 30 31 32# For background on "segments" and "segment ids", see: 33# https://www.tensorflow.org/api_docs/python/tf/math#Segmentation 34@tf_export("ragged.row_splits_to_segment_ids") 35@dispatch.add_dispatch_support 36def row_splits_to_segment_ids(splits, name=None, out_type=None): 37 """Generates the segmentation corresponding to a RaggedTensor `row_splits`. 38 39 Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if 40 `splits[j] <= i < splits[j+1]`. Example: 41 42 >>> print(tf.ragged.row_splits_to_segment_ids([0, 3, 3, 5, 6, 9])) 43 tf.Tensor([0 0 0 2 2 3 4 4 4], shape=(9,), dtype=int64) 44 45 Args: 46 splits: A sorted 1-D integer Tensor. `splits[0]` must be zero. 47 name: A name prefix for the returned tensor (optional). 48 out_type: The dtype for the return value. Defaults to `splits.dtype`, 49 or `tf.int64` if `splits` does not have a dtype. 50 51 Returns: 52 A sorted 1-D integer Tensor, with `shape=[splits[-1]]` 53 54 Raises: 55 ValueError: If `splits` is invalid. 56 """ 57 with ops.name_scope(name, "RaggedSplitsToSegmentIds", [splits]) as name: 58 splits = ops.convert_to_tensor( 59 splits, name="splits", 60 preferred_dtype=dtypes.int64) 61 if splits.dtype not in (dtypes.int32, dtypes.int64): 62 raise ValueError("splits must have dtype int32 or int64") 63 splits.shape.assert_has_rank(1) 64 if tensor_shape.dimension_value(splits.shape[0]) == 0: 65 raise ValueError("Invalid row_splits: []") 66 if out_type is None: 67 out_type = splits.dtype 68 else: 69 out_type = dtypes.as_dtype(out_type) 70 row_lengths = splits[1:] - splits[:-1] 71 nrows = array_ops.shape(splits, out_type=out_type)[-1] - 1 72 indices = math_ops.range(nrows) 73 return ragged_util.repeat(indices, repeats=row_lengths, axis=0) 74 75 76# For background on "segments" and "segment ids", see: 77# https://www.tensorflow.org/api_docs/python/tf/math#Segmentation 78@tf_export("ragged.segment_ids_to_row_splits") 79@dispatch.add_dispatch_support 80def segment_ids_to_row_splits(segment_ids, num_segments=None, 81 out_type=None, name=None): 82 """Generates the RaggedTensor `row_splits` corresponding to a segmentation. 83 84 Returns an integer vector `splits`, where `splits[0] = 0` and 85 `splits[i] = splits[i-1] + count(segment_ids==i)`. Example: 86 87 >>> print(tf.ragged.segment_ids_to_row_splits([0, 0, 0, 2, 2, 3, 4, 4, 4])) 88 tf.Tensor([0 3 3 5 6 9], shape=(6,), dtype=int64) 89 90 Args: 91 segment_ids: A 1-D integer Tensor. 92 num_segments: A scalar integer indicating the number of segments. Defaults 93 to `max(segment_ids) + 1` (or zero if `segment_ids` is empty). 94 out_type: The dtype for the return value. Defaults to `segment_ids.dtype`, 95 or `tf.int64` if `segment_ids` does not have a dtype. 96 name: A name prefix for the returned tensor (optional). 97 98 Returns: 99 A sorted 1-D integer Tensor, with `shape=[num_segments + 1]`. 100 """ 101 # Local import bincount_ops to avoid import-cycle. 102 from tensorflow.python.ops import bincount_ops # pylint: disable=g-import-not-at-top 103 if out_type is None: 104 if isinstance(segment_ids, ops.Tensor): 105 out_type = segment_ids.dtype 106 elif isinstance(num_segments, ops.Tensor): 107 out_type = num_segments.dtype 108 else: 109 out_type = dtypes.int64 110 else: 111 out_type = dtypes.as_dtype(out_type) 112 with ops.name_scope(name, "SegmentIdsToRaggedSplits", [segment_ids]) as name: 113 # Note: we cast int64 tensors to int32, since bincount currently only 114 # supports int32 inputs. 115 segment_ids = ragged_util.convert_to_int_tensor(segment_ids, "segment_ids", 116 dtype=dtypes.int32) 117 segment_ids.shape.assert_has_rank(1) 118 if num_segments is not None: 119 num_segments = ragged_util.convert_to_int_tensor(num_segments, 120 "num_segments", 121 dtype=dtypes.int32) 122 num_segments.shape.assert_has_rank(0) 123 124 row_lengths = bincount_ops.bincount( 125 segment_ids, 126 minlength=num_segments, 127 maxlength=num_segments, 128 dtype=out_type) 129 splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0) 130 131 # Update shape information, if possible. 132 if num_segments is not None: 133 const_num_segments = tensor_util.constant_value(num_segments) 134 if const_num_segments is not None: 135 splits.set_shape(tensor_shape.TensorShape([const_num_segments + 1])) 136 137 return splits 138