Returns new optionals lists that have matching signatures. This is done by mirroring each list in the other using none optionals. There is no merging of like optionals. Args: branch_graphs: `list` of `FuncGraph`. branch_optionals: `list` of `list`s of optional `Tensor`s from other
(branch_graphs, branch_optionals)
| 464 | |
| 465 | |
| 466 | def _make_intermediates_match(branch_graphs, branch_optionals): |
| 467 | """Returns new optionals lists that have matching signatures. |
| 468 | |
| 469 | This is done by mirroring each list in the other using none optionals. |
| 470 | There is no merging of like optionals. |
| 471 | |
| 472 | Args: |
| 473 | branch_graphs: `list` of `FuncGraph`. |
| 474 | branch_optionals: `list` of `list`s of optional `Tensor`s from other |
| 475 | branch_graphs |
| 476 | |
| 477 | Returns: |
| 478 | A `list` of `list`s of `Tensor`s for each branch_graph. Each list has the |
| 479 | same number of `Tensor`s, all of which will be optionals of the same |
| 480 | shape/type. |
| 481 | """ |
| 482 | new_branch_optionals = [] |
| 483 | # Since the intermediates are optionals with dtype variant, we only need |
| 484 | # enough room for the longest list of intermediates. |
| 485 | intermediates_size = max(len(o) for o in branch_optionals) |
| 486 | for i, branch_graph in enumerate(branch_graphs): |
| 487 | other_optionals = _create_none_optionals( |
| 488 | branch_graph, intermediates_size - len(branch_optionals[i])) |
| 489 | new_branch_optionals.append(branch_optionals[i] + other_optionals) |
| 490 | return new_branch_optionals |
| 491 | |
| 492 | |
| 493 | def _make_intermediates_match_xla(branch_graphs, branch_intermediates): |
no test coverage detected