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

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

        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg0) {table_ids = [1, 2], device_ordinal = 1} : (tensor<!tf_type.string>) -> ()
        // expected-remark@above {{ID: 0}}
        // expected-remark@above {{Successors: {1}}}
        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg0) {table_ids = [1, 2], device_ordinal = 1} : (tensor<!tf_type.string>) -> ()
        // expected-remark@above {{ID: 1}}

    // CHECK: %0 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
    // CHECK: %1:4 = "tf._XlaRecvAtHost"(%0) <{device_ordinal = 0 : i64, device_type = "TPU", key = "host_compute_channel_0_args"}> : (tensor<3x!tf_type.string>) -> (tensor<3360x8xi32>, tensor<3xi32>, tensor<i32>, tensor<2xi32>)

                           "_XlaRecvAtHost", opts.op_registry());
  node_builder.Input(std::move(key_input));
  return opts.WithAttr("Toutputs", dtypes)
      .WithAttr("key", key)
      .WithAttr("device_ordinal", 0)
      .WithAttr("_encapsulate", cluster)
      .WithAttr("_outside", oc_cluster)
      .FinalizeBuilder(&node_builder);
}

Node* SendFromHost(ops::NodeOut key_input, const string& cluster,

// Helper function to get an absolute device string, combining device and
// ordinal attribute values.
std::string GetAbsDeviceStr(Operation* op, uint64_t device_ordinal) {
  std::string device_ordinal_str = std::to_string(device_ordinal);
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  if (!device_attr || device_attr.getValue().empty()) return device_ordinal_str;

      %1 = "tf_device.launch"() ({
        %2 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
        %3 = "tf._XlaRecvAtHost"(%2) {_xla_has_host_transfer = true, device_ordinal = 0 : i64, key = "host_compute_channel_0_args"} : (tensor<3x!tf_type.string>) -> tensor<?xi32>
        %4 = "tf.B"(%3) : (tensor<?xi32>) -> tensor<?x!tf_type.string>
        tf_device.return %4 : tensor<?x!tf_type.string>