def _ragged_tile_axis(rt_input, axis, repeats, row_splits_dtype): """Tile a dimension of a RaggedTensor to match a ragged shape.""" assert axis > 0 # Outermost dimension may not be ragged. if not ragged_tensor.is_ragged(rt_input): rt_input = ragged_tensor.RaggedTensor.from_tensor( rt_input, ragged_rank=1, row_splits_dtype=row_splits_dtype) if axis > 1: return rt_input.with_values( _ragged_tile_axis(rt_input.values, axis - 1, repeats, row_splits_dtype)) else: src_row_splits = rt_input.nested_row_splits src_row_lengths = rt_input.nested_row_lengths() splits = src_row_splits[0] dst_row_lengths = [repeats] for i in range(1, len(src_row_lengths)): dst_row_lengths.append( ragged_util.repeat_ranges(src_row_lengths[i], splits, repeats)) splits = array_ops.gather(src_row_splits[i], splits) dst_values = ragged_util.repeat_ranges(rt_input.flat_values, splits, repeats) return ragged_tensor.RaggedTensor.from_nested_row_lengths( dst_values, dst_row_lengths, validate=False)
def _ragged_tile_axis(rt_input, axis, repeats, row_splits_dtype): """Tile a dimension of a RaggedTensor to match a ragged shape.""" assert axis > 0 # Outermost dimension may not be ragged. if not ragged_tensor.is_ragged(rt_input): rt_input = ragged_tensor.RaggedTensor.from_tensor( rt_input, ragged_rank=1, row_splits_dtype=row_splits_dtype) if axis > 1: return rt_input.with_values( _ragged_tile_axis(rt_input.values, axis - 1, repeats, row_splits_dtype)) else: src_row_splits = rt_input.nested_row_splits src_row_lengths = rt_input.nested_row_lengths() splits = src_row_splits[0] dst_row_lengths = [repeats] for i in range(1, len(src_row_lengths)): dst_row_lengths.append( ragged_util.repeat_ranges(src_row_lengths[i], splits, repeats)) splits = array_ops.gather(src_row_splits[i], splits) dst_values = ragged_util.repeat_ranges(rt_input.flat_values, splits, repeats) return ragged_tensor.RaggedTensor.from_nested_row_lengths( dst_values, dst_row_lengths, validate=False)
def _tile_ragged_values(rt_input, multiples, const_multiples=None): """Builds flat_values tensor for a tiled `RaggedTensor`. Returns a tensor that repeats the values in `rt_input.flat_values` in the appropriate pattern to construct a `RaggedTensor` that tiles `rt_input` as specified by `multiples`. Args: rt_input: The `RaggedTensor` whose values should be repeated. multiples: A 1-D integer `tensor`, indicating how many times each dimension should be repeated. const_multiples: Optional constant value for multiples. Used to skip tiling dimensions where `multiples=1`. Returns: A `Tensor` with the same type and rank as `rt_input.flat_values`. #### Example: ```python >>> rt = tf.ragged.constant([[1, 2], [3]]) >>> _tile_ragged_values(rt, [3, 2]) [1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3] ``` """ ragged_rank = rt_input.ragged_rank nested_splits = rt_input.nested_row_splits # Pointers to the values in `rt_input.flat_values`. inner_value_ids = math_ops.range(nested_splits[-1][-1]) # For each ragged dimension (working from the innermost to outermost), # expand `inner_value_ids` as necessary to tile that dimension. prev_splits = None for axis in range(ragged_rank, 0, -1): # Ragged splits for this dimension. splits = nested_splits[axis - 1] # Adjust splits so they point into `inner_value_ids` (instead of just # pointing into the next dimension's values). if prev_splits is not None: # Not the first pass through the loop. splits = array_ops.gather(prev_splits * multiples[axis + 1], splits) # Repeat each element in this ragged dimension `multiples[axis]` times. if const_multiples is None or const_multiples[axis] != 1: inner_value_ids = ragged_util.repeat_ranges( inner_value_ids, splits, multiples[axis]) prev_splits = splits # Gather the tiled inner values. ragged_tiled_values = array_ops.gather(rt_input.flat_values, inner_value_ids) # Tile the flat_values for the uniform dimensions (i.e., for `axis=0` plus # `axis=range(ragged_rank, rank)`). inner_repeats = array_ops.concat( [multiples[:1], multiples[ragged_rank + 1:]], axis=0) return array_ops.tile(ragged_tiled_values, inner_repeats)
def _broadcast_uniform_partitioned_dimension(self, axis, lengths): """Broadcasts the partitioned dimension `axis` to match `lengths`.""" axis_dim_size = self.dimension_size(axis) partitioned_sizes = list(self._partitioned_dim_sizes[:axis]) if lengths.shape.ndims == 0: lengths = array_ops.where( math_ops.equal(axis_dim_size, 1), lengths, axis_dim_size) repeats = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, 1) splits = array_ops.stack([0, self.num_slices_in_dimension(axis)]) else: splits = math_ops.range( array_ops.size(lengths, out_type=self.dim_size_dtype) + 1) repeats = lengths partitioned_sizes.append(lengths) for dim_size in self._partitioned_dim_sizes[axis + 1:]: if dim_size.shape.ndims == 0: partitioned_sizes.append(dim_size) splits *= dim_size else: partitioned_sizes.append( ragged_util.repeat_ranges(dim_size, splits, repeats)) splits = array_ops.gather( ragged_util.lengths_to_splits(dim_size), splits) inner_sizes = self._inner_dim_sizes return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes)
def _broadcast_uniform_partitioned_dimension(self, axis, lengths): """Broadcasts the partitioned dimension `axis` to match `lengths`.""" axis_dim_size = self.dimension_size(axis) partitioned_sizes = list(self._partitioned_dim_sizes[:axis]) if lengths.shape.ndims == 0: lengths = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, axis_dim_size) repeats = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, 1) splits = array_ops.stack([0, self.num_slices_in_dimension(axis)]) else: splits = math_ops.range( array_ops.size(lengths, out_type=self.dim_size_dtype) + 1) repeats = lengths partitioned_sizes.append(lengths) for dim_size in self._partitioned_dim_sizes[axis + 1:]: if dim_size.shape.ndims == 0: partitioned_sizes.append(dim_size) splits *= dim_size else: partitioned_sizes.append( ragged_util.repeat_ranges(dim_size, splits, repeats)) splits = array_ops.gather( ragged_util.lengths_to_splits(dim_size), splits) inner_sizes = self._inner_dim_sizes return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes)
def _tile_ragged_values(rt_input, multiples, const_multiples=None): """Builds flat_values tensor for a tiled `RaggedTensor`. Returns a tensor that repeats the values in `rt_input.flat_values` in the appropriate pattern to construct a `RaggedTensor` that tiles `rt_input` as specified by `multiples`. Args: rt_input: The `RaggedTensor` whose values should be repeated. multiples: A 1-D integer `tensor`, indicating how many times each dimension should be repeated. const_multiples: Optional constant value for multiples. Used to skip tiling dimensions where `multiples=1`. Returns: A `Tensor` with the same type and rank as `rt_input.flat_values`. #### Example: ```python >>> rt = tf.ragged.constant([[1, 2], [3]]) >>> _tile_ragged_values(rt, [3, 2]) [1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3] ``` """ ragged_rank = rt_input.ragged_rank nested_splits = rt_input.nested_row_splits # Pointers to the values in `rt_input.flat_values`. inner_value_ids = math_ops.range(nested_splits[-1][-1]) # For each ragged dimension (working from the innermost to outermost), # expand `inner_value_ids` as necessary to tile that dimension. prev_splits = None for axis in range(ragged_rank, 0, -1): # Ragged splits for this dimension. splits = nested_splits[axis - 1] # Adjust splits so they point into `inner_value_ids` (instead of just # pointing into the next dimension's values). if prev_splits is not None: # Not the first pass through the loop. splits = array_ops.gather(prev_splits * multiples[axis + 1], splits) # Repeat each element in this ragged dimension `multiples[axis]` times. if const_multiples is None or const_multiples[axis] != 1: inner_value_ids = ragged_util.repeat_ranges(inner_value_ids, splits, multiples[axis]) prev_splits = splits # Gather the tiled inner values. ragged_tiled_values = array_ops.gather(rt_input.flat_values, inner_value_ids) # Tile the flat_values for the uniform dimensions (i.e., for `axis=0` plus # `axis=range(ragged_rank, rank)`). inner_repeats = array_ops.concat([multiples[:1], multiples[ragged_rank + 1:]], axis=0) return array_ops.tile(ragged_tiled_values, inner_repeats)
def _tile_ragged_splits(rt_input, multiples, const_multiples=None): """Builds nested_split tensors for a tiled `RaggedTensor`. Returns a list of split tensors that can be used to construct the `RaggedTensor` that tiles `rt_input` as specified by `multiples`. Args: rt_input: The `RaggedTensor` that is being tiled. multiples: A 1-D integer `tensor`, indicating how many times each dimension should be repeated. const_multiples: Optional constant value for multiples. Used to skip tiling dimensions where `multiples=1`. Returns: A list of 1-D integer `Tensor`s (one for each ragged dimension in `rt_input`). #### Example: >>> rt = tf.ragged.constant([[1, 2], [3]]) >>> _tile_ragged_splits(rt, [3, 2]) [<tf.Tensor: shape=(7,), dtype=int64, numpy=array([ 0, 4, 6, 10, 12, 16, 18])>] """ ragged_rank = rt_input.ragged_rank nested_splits = rt_input.nested_row_splits # projected_splits[src_axis, dst_axis] contains the split points that divide # the rows from src_axis in the list of dst_axis values. E.g., # projected_splits[i, i] = nested_splits[i], and # projected_splits[i, i+1] = gather(nested_splits[i+1], nested_splits[i]). projected_splits = [{i: nested_splits[i]} for i in range(ragged_rank)] for src_axis in range(ragged_rank): for dst_axis in range(src_axis + 1, ragged_rank - 1): projected_splits[src_axis][dst_axis] = array_ops.gather( nested_splits[dst_axis], projected_splits[src_axis][dst_axis - 1]) # For each ragged dimension: nested_splits[axis] -> result_splits[axis]. result_splits = [] for axis in range(ragged_rank): # Get the length of each row for the input tensor for this dimension. input_lengths = nested_splits[axis][1:] - nested_splits[axis][:-1] # Multiply those lengths by the `multiples` of dimension axis+1, since # each value will be repeated that number of times. output_lengths = input_lengths * multiples[axis + 1] # Repeat ranges of the row lengths as necessary for them to be tiled in # each ragged dimension `d < axis`. (Start with dimension d=axis-1, and # work our way up to dimension d=0.) repeats = 1 for d in range(axis - 1, -1, -1): if const_multiples is None or const_multiples[d + 1] != 1: splits = projected_splits[d][axis - 1] * repeats output_lengths = ragged_util.repeat_ranges(output_lengths, splits, multiples[d + 1]) repeats *= multiples[d + 1] # Tile splits for the outermost (uniform) dimension. output_lengths = array_ops.tile(output_lengths, multiples[:1]) # Convert to splits. result_splits.append(ragged_util.lengths_to_splits(output_lengths)) return result_splits
def _tile_ragged_splits(rt_input, multiples, const_multiples=None): """Builds nested_split tensors for a tiled `RaggedTensor`. Returns a list of split tensors that can be used to construct the `RaggedTensor` that tiles `rt_input` as specified by `multiples`. Args: rt_input: The `RaggedTensor` that is being tiled. multiples: A 1-D integer `tensor`, indicating how many times each dimension should be repeated. const_multiples: Optional constant value for multiples. Used to skip tiling dimensions where `multiples=1`. Returns: A list of 1-D integer `Tensor`s (one for each ragged dimension in `rt_input`). #### Example: ```python >>> rt = tf.ragged.constant([[1, 2], [3]]) >>> _tile_ragged_splits(rt, [3, 2]) [0, 4, 6, 10, 12, 16, 18] ``` """ ragged_rank = rt_input.ragged_rank nested_splits = rt_input.nested_row_splits # projected_splits[src_axis, dst_axis] contains the split points that divide # the rows from src_axis in the list of dst_axis values. E.g., # projected_splits[i, i] = nested_splits[i], and # projected_splits[i, i+1] = gather(nested_splits[i+1], nested_splits[i]). projected_splits = [{i: nested_splits[i]} for i in range(ragged_rank)] for src_axis in range(ragged_rank): for dst_axis in range(src_axis + 1, ragged_rank - 1): projected_splits[src_axis][dst_axis] = array_ops.gather( nested_splits[dst_axis], projected_splits[src_axis][dst_axis - 1]) # For each ragged dimension: nested_splits[axis] -> result_splits[axis]. result_splits = [] for axis in range(ragged_rank): # Get the length of each row for the input tensor for this dimension. input_lengths = nested_splits[axis][1:] - nested_splits[axis][:-1] # Multiply those lengths by the `multiples` of dimension axis+1, since # each value will be repeated that number of times. output_lengths = input_lengths * multiples[axis + 1] # Repeat ranges of the row lengths as necessary for them to be tiled in # each ragged dimension `d < axis`. (Start with dimension d=axis-1, and # work our way up to dimension d=0.) repeats = 1 for d in range(axis - 1, -1, -1): if const_multiples is None or const_multiples[d + 1] != 1: splits = projected_splits[d][axis - 1] * repeats output_lengths = ragged_util.repeat_ranges(output_lengths, splits, multiples[d + 1]) repeats *= multiples[d + 1] # Tile splits for the outermost (uniform) dimension. output_lengths = array_ops.tile(output_lengths, multiples[:1]) # Convert to splits. result_splits.append(ragged_util.lengths_to_splits(output_lengths)) return result_splits