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 ",

  std::unique_ptr<FunctionBody> host_fbody;
  AttrValue device_ordinal_temp_value;
  device_ordinal_temp_value.set_i(0);
  protobuf::Map<string, AttrValue> host_func_attrs;
  host_func_attrs["_device_ordinal"] = device_ordinal_temp_value;
  TF_CHECK_OK(FunctionDefToBodyHelper(
      *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld, &host_fbody));
  Graph *host_graph = host_fbody->graph;

  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,

    // CHECK-DAG:        %[[DEVICE_ORDINAL:.+]] = "tf._TPUDeviceOrdinalPlaceholder"
    // CHECK:            %[[RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]])
    // CHECK:            %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[RECV_OUTPUT]])
    // CHECK:            "tf._XlaSendFromHostV2"(%[[B_OUTPUT]], %[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]])

  XlaCompiler::Options options;
  options.device_ordinal = 0;
  options.alias_passthrough_params = true;
  options.detailed_logging = true;
  XlaCompiler::CompilationResult result;
  xla::Shape xla_output_shape;
  result.xla_output_shape = xla_output_shape;

  auto build_option =
      GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);

  EXPECT_EQ(build_option.device_ordinal(), 0);

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?="

  // CHECK-SAME:   device_ordinal = 0
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 1
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 0
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 1
}

// -----