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`, indic
(rt_input, multiples, const_multiples=None)
| 309 | |
| 310 | |
| 311 | def _tile_ragged_splits(rt_input, multiples, const_multiples=None): |
| 312 | """Builds nested_split tensors for a tiled `RaggedTensor`. |
| 313 | |
| 314 | Returns a list of split tensors that can be used to construct the |
| 315 | `RaggedTensor` that tiles `rt_input` as specified by `multiples`. |
| 316 | |
| 317 | Args: |
| 318 | rt_input: The `RaggedTensor` that is being tiled. |
| 319 | multiples: A 1-D integer `tensor`, indicating how many times each dimension |
| 320 | should be repeated. |
| 321 | const_multiples: Optional constant value for multiples. Used to skip tiling |
| 322 | dimensions where `multiples=1`. |
| 323 | |
| 324 | Returns: |
| 325 | A list of 1-D integer `Tensor`s (one for each ragged dimension in |
| 326 | `rt_input`). |
| 327 | |
| 328 | #### Example: |
| 329 | ```python |
| 330 | >>> rt = tf.ragged.constant([[1, 2], [3]]) |
| 331 | >>> _tile_ragged_splits(rt, [3, 2]) |
| 332 | [0, 4, 6, 10, 12, 16, 18] |
| 333 | ``` |
| 334 | """ |
| 335 | ragged_rank = rt_input.ragged_rank |
| 336 | nested_splits = rt_input.nested_row_splits |
| 337 | |
| 338 | # projected_splits[src_axis, dst_axis] contains the split points that divide |
| 339 | # the rows from src_axis in the list of dst_axis values. E.g., |
| 340 | # projected_splits[i, i] = nested_splits[i], and |
| 341 | # projected_splits[i, i+1] = gather(nested_splits[i+1], nested_splits[i]). |
| 342 | projected_splits = [{i: nested_splits[i]} for i in range(ragged_rank)] |
| 343 | for src_axis in range(ragged_rank): |
| 344 | for dst_axis in range(src_axis + 1, ragged_rank - 1): |
| 345 | projected_splits[src_axis][dst_axis] = array_ops.gather( |
| 346 | nested_splits[dst_axis], |
| 347 | projected_splits[src_axis][dst_axis - 1]) |
| 348 | |
| 349 | # For each ragged dimension: nested_splits[axis] -> result_splits[axis]. |
| 350 | result_splits = [] |
| 351 | for axis in range(ragged_rank): |
| 352 | # Get the length of each row for the input tensor for this dimension. |
| 353 | input_lengths = nested_splits[axis][1:] - nested_splits[axis][:-1] |
| 354 | |
| 355 | # Multiply those lengths by the `multiples` of dimension axis+1, since |
| 356 | # each value will be repeated that number of times. |
| 357 | output_lengths = input_lengths * multiples[axis + 1] |
| 358 | |
| 359 | # Repeat ranges of the row lengths as necessary for them to be tiled in |
| 360 | # each ragged dimension `d < axis`. (Start with dimension d=axis-1, and |
| 361 | # work our way up to dimension d=0.) |
| 362 | repeats = 1 |
| 363 | for d in range(axis - 1, -1, -1): |
| 364 | if const_multiples is None or const_multiples[d + 1] != 1: |
| 365 | splits = projected_splits[d][axis - 1] * repeats |
| 366 | output_lengths = ragged_util.repeat_ranges(output_lengths, splits, |
| 367 | multiples[d + 1]) |
| 368 | repeats *= multiples[d + 1] |