Python wrapper for an XLA Computation. A Computation can be compiled to form an Executable, or used as a subcomputation in ComputationBuilder methods.
| 524 | |
| 525 | |
| 526 | class Computation(object): |
| 527 | """Python wrapper for an XLA Computation. |
| 528 | |
| 529 | A Computation can be compiled to form an Executable, or used as a |
| 530 | subcomputation in ComputationBuilder methods. |
| 531 | """ |
| 532 | |
| 533 | def __init__(self, c_computation, backend=None): |
| 534 | self._c_computation = c_computation |
| 535 | # The backend argument is deprecated. Pass a backend to Compile() instead. |
| 536 | self._backend = backend |
| 537 | |
| 538 | @property |
| 539 | def computation(self): |
| 540 | return self._c_computation |
| 541 | |
| 542 | def GetSerializedProto(self): |
| 543 | """Gets the serialized HloModuleProto proto object in this computation. |
| 544 | |
| 545 | Returns: |
| 546 | A string containing a serialized HloModuleProto proto containing the |
| 547 | computation and its dependencies. |
| 548 | """ |
| 549 | return self.computation.GetSerializedProto() |
| 550 | |
| 551 | def GetHloText(self): |
| 552 | """Get the textual HLO representation of this computation. |
| 553 | |
| 554 | Returns: |
| 555 | A string containing the textual HLO. |
| 556 | """ |
| 557 | return self.computation.GetHloText() |
| 558 | |
| 559 | def GetHloDotGraph(self): |
| 560 | """Get a Graphviz Dot representation of this computation. |
| 561 | |
| 562 | Returns: |
| 563 | A string containing the graphviz dot graph. |
| 564 | """ |
| 565 | return self.computation.GetHloDotGraph() |
| 566 | |
| 567 | def Compile(self, argument_shapes=None, compile_options=None, backend=None): |
| 568 | """Compiles a computation. |
| 569 | |
| 570 | Computations are the result of a "ComputationBuild'ing" process. |
| 571 | |
| 572 | Arguments: |
| 573 | argument_shapes: Deprecated. Use compile_options.argument_layouts instead. |
| 574 | compile_options: options to use for compilation, includes an optional laid |
| 575 | out result shape for the computation. |
| 576 | backend: a `Backend` for which an executable should be generated. |
| 577 | |
| 578 | Returns: |
| 579 | A Executable instance. |
| 580 | """ |
| 581 | backend = backend or self._backend or get_local_backend() |
| 582 | |
| 583 | compile_options = compile_options or CompileOptions() |