}

mlir::LogicalResult GetDeviceOrdinalFromDeviceString(mlir::Location loc,
                                                     llvm::StringRef device,
                                                     int64_t* device_ordinal) {
  DeviceNameUtils::ParsedName parsed_name;
  if (!DeviceNameUtils::ParseFullName(
          absl::string_view(device.data(), device.size()), &parsed_name))

  // 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

      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{

  }
  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{

  xla::ExecutableBuildOptions build_options;
  if (result.collective_info) {
    build_options.set_num_replicas(result.collective_info->group_size);
  }
  build_options.set_device_ordinal(
      options.device_ordinal != -1 ? options.device_ordinal
                                   : local_client->default_device_ordinal());
  build_options.set_result_layout(result.xla_output_shape);
  build_options.set_device_allocator(options.device_allocator.get());

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