Method RunAsync

tensorflow/compiler/xla/client/local_client.cc:231–258

}

StatusOr<ScopedShapedBuffer> LocalExecutable::RunAsync(
    const absl::Span<const ShapedBuffer* const> arguments,
    ExecutableRunOptions run_options) {
  std::vector<const Shape*> argument_shapes;
  argument_shapes.reserve(arguments.size());
  for (const ShapedBuffer* const arg : arguments) {
    argument_shapes.push_back(&arg->on_host_shape());
  }
  TF_ASSIGN_OR_RETURN(auto options_and_stream,
                      RunHelper(argument_shapes, run_options));
  se::Stream* stream = run_options.stream();

  std::shared_ptr<HloSnapshot> snapshot;
  if (executable_->dumping_snapshot()) {
    snapshot = DumpArguments(backend_, executable_.get(), arguments, stream);
  }

  TF_ASSIGN_OR_RETURN(ScopedShapedBuffer outputs,
                      executable_->ExecuteAsyncOnStreamWrapper(
                          &options_and_stream.first, arguments));

  // Transfer the outputs and save the snapshot to disk.
  if (snapshot) {
    DumpOutputsAndSaveSnapshot(backend_, outputs, std::move(snapshot), stream);
  }

  return std::move(outputs);
}

static ShapedBuffer MaybeOwningShapeTreeToShapedBuffer(
    Shape const& on_host_shape, const ShapeTree<MaybeOwningDeviceMemory>& tree,