Validates flat outputs and adds back device assignments. Args: outputs: Output from `computation` inside `xla.compile`. Returns: Tensors and Operations extracted from outputs.
(outputs)
| 430 | |
| 431 | |
| 432 | def _postprocess_flat_outputs(outputs): |
| 433 | """Validates flat outputs and adds back device assignments. |
| 434 | |
| 435 | Args: |
| 436 | outputs: Output from `computation` inside `xla.compile`. |
| 437 | |
| 438 | Returns: |
| 439 | Tensors and Operations extracted from outputs. |
| 440 | """ |
| 441 | # Following code segment is to preserve legacy behavior. Previously we only |
| 442 | # supported flat outputs and thus for consistency it was nice to convert even |
| 443 | # single element into a tuple. But now that we support arbitrary output |
| 444 | # structure, this is no longer necessary. |
| 445 | # TODO(b/121383831): Migrate all legacy use cases and delete this special |
| 446 | # case. |
| 447 | # If the computation returns `None`, make it an empty tuple. |
| 448 | if outputs is None: |
| 449 | outputs = tuple() |
| 450 | # If the computation only returned one value, make it a tuple. |
| 451 | if not isinstance(outputs, collections.Sequence): |
| 452 | outputs = (outputs,) |
| 453 | |
| 454 | # Append `no_op` here so that return value of this function always contains |
| 455 | # at least one op that can trigger XlaLaunch node. |
| 456 | outputs += (control_flow_ops.no_op(),) |
| 457 | try: |
| 458 | outputs = [ |
| 459 | o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) |
| 460 | for o in outputs |
| 461 | ] |
| 462 | except Exception as e: |
| 463 | raise ValueError( |
| 464 | 'XLA computation function return values must all either be Operations' |
| 465 | ' or convertible to Tensors. Got error: "%s"' % str(e)) |
| 466 | |
| 467 | # Separates the returned Operations and Tensors. |
| 468 | output_operations = [o for o in outputs if isinstance(o, ops.Operation)] |
| 469 | output_tensors = [o for o in outputs if not isinstance(o, ops.Operation)] |
| 470 | |
| 471 | if outputs != output_tensors + output_operations: |
| 472 | raise ValueError( |
| 473 | 'XLA computation function must return zero or more Tensor values ' |
| 474 | 'followed by zero or more Operations.') |
| 475 | |
| 476 | new_output_tensors = [] |
| 477 | for t in output_tensors: |
| 478 | with ops.device(t.device if t.device else ''): |
| 479 | new_output_tensors.append(array_ops.identity(t)) |
| 480 | |
| 481 | return new_output_tensors, output_operations |
| 482 | |
| 483 | |
| 484 | def _postprocess_non_flat_outputs(outputs): |
no test coverage detected