n->ClearAttr(attr_name);
        n->AddAttr(attr_name, branch_func);
      }
    } else if (HasNodeAttr(n->def(), "_device_ordinal")) {
      // Function call node containing outside compilation.
      n->ClearAttr("_device_ordinal");
      n->AddAttr("_device_ordinal", device_ordinal_value);
    } else {
      return errors::Internal("Unknown node marked with ",

  int64_t device_ordinal0 = -1;
  mlir::LogicalResult result0 =
      GetDeviceOrdinalFromDeviceString(unknown_loc, tpu0, &device_ordinal0);
  EXPECT_TRUE(mlir::succeeded(result0));
  EXPECT_EQ(device_ordinal0, 0);

  int64_t device_ordinal1 = -1;
  mlir::LogicalResult result1 =
      GetDeviceOrdinalFromDeviceString(unknown_loc, tpu1, &device_ordinal1);
  EXPECT_TRUE(mlir::succeeded(result1));

    : 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_; }

  class Metadata {
   public:
    Metadata(int device_ordinal, se::Platform* platform,
             const DeviceType& device_type,
             std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                 shape_determination_fns,
             PaddedShapeFn padded_shape_fn, bool use_multiple_streams);

    // The index of the device on this host.
    int device_ordinal() const;

    se::Platform* platform() const;

                                Value device_ordinal,
                                StringAttr device_type_attr,
                                llvm::StringRef communication_key) {
  if (device_ordinal)
    return ApplyXlaHostTransferAttr(
        builder.create<mlir::TF::_XlaSendFromHostV2Op>(
            loc, inputs,
            /*dynamic_key=*/compilation_key, device_ordinal,

inline void LogOptions(const XlaCompiler::Options& options) {
  VLOG(2) << "XlaCompiler::Options[device_type=" << options.device_type
          << ",device_ordinal=" << options.device_ordinal
          << ",client=" << options.client << ",flib_def=" << options.flib_def
          << ",graph_def_version=" << options.graph_def_version
          << ",options.shape_determination_fns.layout_preference_fn?="

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

    Options(std::string device_name_prefix, std::string device_name,
            int device_ordinal, std::string compilation_device_name,
            std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                shape_determination_fns)
        : device_name_prefix(device_name_prefix),
          device_name(device_name),
          device_ordinal(device_ordinal),
          compilation_device_name(compilation_device_name),

                                               const string& device_name,
                                               int device_ordinal) {
  return Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
      DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
      absl::StrCat("device: ", device_name, " device"));
}

}  // namespace

                                       int device_ordinal) {
  xla::Shape on_host_shape =
      xla::ShapeUtil::DeviceShapeToHostShape(on_device_shape);
  xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,
                                        client->backend().memory_allocator(),
                                        device_ordinal);
  for (auto& index_to_buffer : shaped_buffer.buffers()) {