Validates non-flat outputs and adds back device assignments. Args: outputs: Output from `computation` inside `xla.compile`. Returns: Tensors extracted from outputs and an empty list because Operations are not allowed in non-flat outputs..
(outputs)
| 482 | |
| 483 | |
| 484 | def _postprocess_non_flat_outputs(outputs): |
| 485 | """Validates non-flat outputs and adds back device assignments. |
| 486 | |
| 487 | Args: |
| 488 | outputs: Output from `computation` inside `xla.compile`. |
| 489 | |
| 490 | Returns: |
| 491 | Tensors extracted from outputs and an empty list because Operations are not |
| 492 | allowed in non-flat outputs.. |
| 493 | """ |
| 494 | # Convert all non-Operation outputs to Tensors. |
| 495 | new_output_tensors = [] |
| 496 | for o in nest.flatten(outputs): |
| 497 | if isinstance(o, ops.Operation): |
| 498 | raise ValueError( |
| 499 | 'xla.compile does not support Operation as return value in non-flat ' |
| 500 | 'output structure. You can set returned Operations as control ' |
| 501 | 'dependencies of returned Tensors so Operations are triggered when ' |
| 502 | 'Tensors are evaluated. Operation found: "%s"' % o.name) |
| 503 | |
| 504 | try: |
| 505 | o = ops.convert_to_tensor(o) |
| 506 | except Exception as e: |
| 507 | raise ValueError( |
| 508 | 'XLA computation function return values must all either be ' |
| 509 | 'Operations or convertible to Tensors. Got error: "%s"' % str(e)) |
| 510 | |
| 511 | # Makes sure even pass-through inputs/outputs are touched in compile |
| 512 | # context by creating an Identity node inside compile context. |
| 513 | with ops.device(o.device if o.device else ''): |
| 514 | new_output_tensors.append(array_ops.identity(o)) |
| 515 | |
| 516 | return new_output_tensors, [] |
| 517 | |
| 518 | |
| 519 | @contextlib.contextmanager |
no test coverage detected