(x)
| 182 | output_is_sequence = nest.is_sequence(dtype) |
| 183 | output_flatten = lambda x: nest.flatten(x) if output_is_sequence else [x] |
| 184 | def output_pack(x): |
| 185 | return (nest.pack_sequence_as(dtype, x) |
| 186 | if output_is_sequence else x[0]) |
| 187 | |
| 188 | elems_flat = input_flatten(elems) |
| 189 |