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