Merge pull request #52633 from pranve/cherrypick-874bda09e6702cd50bac…

…90b453b50bcc65b2769e-on-r2.6 Merge pull request #51715 from yongtang:46909-tf.summary.create_file_…
tensorflow · Oct 25, 2021 · 3c25d95 · 3c25d95
2 parents 64f74eb + 7bbf1df
commit 3c25d95
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/tensorflow/core/kernels/summary_kernels.cc b/tensorflow/core/kernels/summary_kernels.cc
@@ -38,12 +38,22 @@ class CreateSummaryFileWriterOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor* tmp;
     OP_REQUIRES_OK(ctx, ctx->input("logdir", &tmp));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tmp->shape()),
+                errors::InvalidArgument("logdir must be a scalar"));
     const string logdir = tmp->scalar<tstring>()();
     OP_REQUIRES_OK(ctx, ctx->input("max_queue", &tmp));
-    const int32 max_queue = tmp->scalar<int32>()();
+
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tmp->shape()),
+                errors::InvalidArgument("max_queue must be a scalar"));
+    const int32_t max_queue = tmp->scalar<int32>()();
     OP_REQUIRES_OK(ctx, ctx->input("flush_millis", &tmp));
-    const int32 flush_millis = tmp->scalar<int32>()();
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tmp->shape()),
+                errors::InvalidArgument("flush_millis must be a scalar"));
+    const int32_t flush_millis = tmp->scalar<int32>()();
+
     OP_REQUIRES_OK(ctx, ctx->input("filename_suffix", &tmp));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tmp->shape()),
+                errors::InvalidArgument("filename_suffix must be a scalar"));
     const string filename_suffix = tmp->scalar<tstring>()();
 
     core::RefCountPtr<SummaryWriterInterface> s;

diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py
@@ -34,6 +34,7 @@
 from tensorflow.python.client import session
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import meta_graph
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
@@ -685,6 +686,16 @@ def testSharing_withExplicitSummaryFileWriters(self):
     # No more files
     self.assertRaises(StopIteration, lambda: next(event_paths))
 
+  def testSummaryFileWritersInvalidInput(self):
+    # Test case for GitHub issue 46909
+    logdir = self.get_temp_dir()
+    with session.Session() as sess:
+      with self.assertRaises(errors_impl.InvalidArgumentError):
+        writer = summary_ops_v2.create_file_writer(
+            logdir=logdir, flush_millis=[1, 2])
+        sess.run(writer.init())
+        sess.run(writer.flush())
+
 
 class FileWriterCacheTest(test.TestCase):
   """FileWriterCache tests."""