Like `nest.pack_sequence_as` but also replaces flows with TensorArrays.
(structure_with_tas, loop_vars)
| 1106 | |
| 1107 | |
| 1108 | def _pack_sequence_as(structure_with_tas, loop_vars): |
| 1109 | """Like `nest.pack_sequence_as` but also replaces flows with TensorArrays.""" |
| 1110 | |
| 1111 | def flow_to_tensor_array(flow, ta): # pylint: disable=missing-docstring |
| 1112 | return (tensor_array_ops.build_ta_with_new_flow(ta, flow) if isinstance( # pylint: disable=g-long-ternary |
| 1113 | ta, tensor_array_ops.TensorArray) else flow) |
| 1114 | |
| 1115 | flattened_loop_vars = [ |
| 1116 | flow_to_tensor_array(*z) |
| 1117 | for z in zip(nest.flatten(loop_vars, expand_composites=True), |
| 1118 | nest.flatten(structure_with_tas, expand_composites=True)) |
| 1119 | ] |
| 1120 | return nest.pack_sequence_as(structure_with_tas, flattened_loop_vars, |
| 1121 | expand_composites=True) |
| 1122 | |
| 1123 | |
| 1124 | def _tensor_array_to_flow(loop_vars): |
no test coverage detected