// objects.
  XlaComputationLaunchContext(xla::LocalClient* client,
                              se::DeviceMemoryAllocator* xla_allocator,
                              int device_ordinal, bool allocate_xla_tensors,
                              bool use_multiple_streams);

  // Builds a XlaCompiler::Argument vector from the arguments to an XlaLaunch
  // op.

        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

mlir::LogicalResult GetDeviceOrdinalFromDeviceString(mlir::Location loc,
                                                     llvm::StringRef device,
                                                     int64_t* device_ordinal);

}  // namespace tensorflow

// CHECK:      tf_executor.fetch


// Tests tf._TPUDeviceOrdinalPlaceholder ops are replaced with explicit device
// ordinal constant values based on the first TPU core device id.
// CHECK-LABEL: func @device_ordinals
func.func @device_ordinals() {
  tf_executor.graph {
    %0:3 = tf_executor.island {

      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())) {

  XlaDevice::Options options;
  options.platform = platform;
  options.device_name_prefix = name_prefix;
  options.device_name = DEVICE_XLA_CPU;
  options.device_ordinal = 0;
  options.compilation_device_name = DEVICE_CPU_XLA_JIT;
  options.use_multiple_streams = false;
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{

  // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()

  }
  for (int i : *gpu_ids) {
    XlaDevice::Options options;
    options.platform = platform.value();
    options.device_name_prefix = name_prefix;
    options.device_name = DEVICE_XLA_GPU;
    options.device_ordinal = i;
    options.compilation_device_name = DEVICE_GPU_XLA_JIT;
    options.use_multiple_streams = true;
    options.allowed_devices = gpu_ids;
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{

  // is replaced and the managed memory deallocated.
  Status AllocateShapedBuffer(DataType dtype, const xla::Shape& on_device_shape,
                              xla::LocalClient* client, int device_ordinal);

  // Some Tensors can have complex on-device shapes, including tuple shapes. To
  // manage the memory for these tensors a ShapedBuffer may be required.

  // Return true if this XlaTensor contains a ShapedBuffer.