: device_ordinal_(device_ordinal),
      device_type_(device_type),
      platform_(platform),
      shape_determination_fns_(std::move(shape_determination_fns)),
      padded_shape_fn_(std::move(padded_shape_fn)),
      use_multiple_streams_(use_multiple_streams) {}

int XlaDevice::Metadata::device_ordinal() const { return device_ordinal_; }

    xla::LocalClient* client, se::DeviceMemoryAllocator* xla_allocator,
    int device_ordinal, bool allocate_xla_tensors, bool use_multiple_streams)
    : client_(client),
      xla_allocator_(xla_allocator),
      allocate_xla_tensors_(allocate_xla_tensors),
      use_multiple_streams_(use_multiple_streams),
      device_ordinal_(device_ordinal) {
  if (use_multiple_streams_) {

        input->dtype(), shape, dst_xla_context->client(), dst_device_ordinal));

    VLOG(2) << "TpuDeviceToDeviceCopy: src: "
            << src_compute_stream->parent()->device_ordinal() << ", "
            << " dst: " << dst_compute_stream->parent()->device_ordinal()
            << ", "
            << " input buffers: " << xla_input->shaped_buffer().ToString()
            << " output buffers: " << xla_output->shaped_buffer().ToString();

    TF_RET_CHECK(!xla_tensor->has_shaped_buffer());

    TF_RETURN_IF_ERROR(
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal()));

    // The cpu_tensor and literal that we created here hold the data of host
    // tensor in descending layout. The layout could be different from layout in

      auto recv_at_host = rewriter.create<TF::_XlaRecvAtHostOp>(
          func.getLoc(), op.getOperandTypes(), /*dynamic_key=*/dynamic_key,
          op.getSendKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
      for (auto result :
           llvm::zip(cloned_func.getArguments(), recv_at_host->getResults())) {

    xla::LocalClient* client, se::DeviceMemoryAllocator* allocator) {
  se::Stream* stream = GetStream(ctx);
  int device_ordinal = stream ? stream->parent()->device_ordinal()
                              : client->default_device_ordinal();
  XlaComputationLaunchContext launch_context(
      client, allocator, device_ordinal,
      /*allocate_xla_tensors=*/platform_info.is_on_xla_device(),

          %2 = "tf._XlaRecvAtHost"(%1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_args"} : (tensor<3x!tf_type.string>) -> tensor<f32>
          %3 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>
          "tf._XlaSendFromHost"(%3, %1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_retvals"} : (tensor<f32>, tensor<3x!tf_type.string>) -> ()

  Output x = ops::Placeholder(root.WithOpName("x"), DT_FLOAT);
  NameAttrList f_name_attr;
  f_name_attr.set_name("XTimesTwoFloat");
  ops::TPUPartitionedCall f(root.WithOpName("f"), {x}, /*device_ordinal=*/0,
                            {DT_FLOAT}, f_name_attr);

  TF_ASSERT_OK(root.ToGraph(&graph));
  EXPECT_TRUE(IsInferenceGraph(graph, /*function_library=*/nullptr));
}