Tile a dimension of a RaggedTensor to match a ragged shape.
(rt_input, axis, repeats, row_splits_dtype)
| 599 | |
| 600 | |
| 601 | def _ragged_tile_axis(rt_input, axis, repeats, row_splits_dtype): |
| 602 | """Tile a dimension of a RaggedTensor to match a ragged shape.""" |
| 603 | assert axis > 0 # Outermost dimension may not be ragged. |
| 604 | |
| 605 | if not ragged_tensor.is_ragged(rt_input): |
| 606 | rt_input = ragged_tensor.RaggedTensor.from_tensor( |
| 607 | rt_input, ragged_rank=1, row_splits_dtype=row_splits_dtype) |
| 608 | |
| 609 | if axis > 1: |
| 610 | return rt_input.with_values( |
| 611 | _ragged_tile_axis(rt_input.values, axis - 1, repeats, |
| 612 | row_splits_dtype)) |
| 613 | else: |
| 614 | src_row_splits = rt_input.nested_row_splits |
| 615 | src_row_lengths = rt_input.nested_row_lengths() |
| 616 | splits = src_row_splits[0] |
| 617 | |
| 618 | dst_row_lengths = [repeats] |
| 619 | for i in range(1, len(src_row_lengths)): |
| 620 | dst_row_lengths.append( |
| 621 | ragged_util.repeat_ranges(src_row_lengths[i], splits, repeats)) |
| 622 | splits = array_ops.gather(src_row_splits[i], splits) |
| 623 | dst_values = ragged_util.repeat_ranges(rt_input.flat_values, splits, |
| 624 | repeats) |
| 625 | return ragged_tensor.RaggedTensor.from_nested_row_lengths( |
| 626 | dst_values, dst_row_lengths, validate=False) |
no test coverage detected