From 39dbb1f792e30ca0d4791432bd77f9582b158892 Mon Sep 17 00:00:00 2001
From: Jim Lin <jimlintw@google.com>
Date: Wed, 8 May 2024 16:43:58 -0700
Subject: [PATCH] #tf-data implements string type AddScalar

PiperOrigin-RevId: 631965339
---
 tensorflow/core/framework/dataset.h | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index effa997eca60f9..1237fbdec449da 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -227,6 +227,16 @@ class GraphDefBuilderWrapper {
     return OkStatus();
   }
 
+  absl::Status AddScalar(const std::string& val, Node** output) {
+    Tensor val_t = Tensor(DataTypeToEnum<tstring>::v(), TensorShape({}));
+    val_t.scalar<tstring>()() = val;
+    AddTensorInternal(val_t, output);
+    if (*output == nullptr) {
+      return absl::InternalError("AddScalar: Failed to build Const op.");
+    }
+    return absl::OkStatus();
+  }
+
   // Adds a Const node with vector value to the Graph.
   // `*output` contains a pointer to the output `Node`. It is guaranteed to be
   // non-null if the method returns with an OK status.