Add new data transfer support.

PiperOrigin-RevId: 361002139 Change-Id: I779bfffeb8659e2d31381cdffdacc2cda591081b
2021-03-04 14:17:06 -08:00 · 2021-03-04 14:17:06 -08:00 · fa12754826
commit fa12754826
parent 6a2d982fa9
1 changed files with 2 additions and 0 deletions
--- a/tensorflow/compiler/xla/service/hlo_runner_interface.h
+++ b/tensorflow/compiler/xla/service/hlo_runner_interface.h
@ -169,6 +169,8 @@ class HloRunnerInterface {
      std::function<int64(int64)> argument_count_provider,
      std::function<const Literal*(int64, int64)> argument_provider,
      const ReplicatedExecuteOptions& options) = 0;
+
+  typedef std::function<Shape(const Shape&)> TpuShapeRepresentationFn;
 };

 }  // namespace xla