commit 8cef4cda26e08256b6698e942820d9a3ac1bcc94
author Prateek Gupta <prateek@polymagelabs.com> Thu Nov 12 11:17:45 2020 +0530
committer Prateek Gupta <prateek@polymagelabs.com> Fri Dec 11 11:52:43 2020 +0530
tree 794e195b90bb600c4f4066af5dcaebb153b76059
parent 12217d1a5e266b5ac9214fca755e67bd349d0b5f

Add minor checks for data_format and padding value

tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc
index 5c6496e..edf90f1 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc
@@ -1467,9 +1467,15 @@
   ArrayRef<Attribute> dilations = op.dilations().getValue();
 
   tensorflow::TensorFormat format;
-  FormatFromString(data_format.str(), &format);
+  auto data_format_is_valid = FormatFromString(data_format.str(), &format);
+  if (!data_format_is_valid) {
+    return emitOptionalError(location, "Invalid data format provided");
+  }
   tensorflow::Padding padding;
-  GetPaddingFromString(paddings.str(), &padding);
+  auto padding_is_valid = GetPaddingFromString(paddings.str(), &padding);
+  if (!padding_is_valid.ok()) {
+    return emitOptionalError(location, "Invalid padding format provided");
+  }
   auto get_int = [](Attribute attr) {
     return attr.template cast<IntegerAttr>().getInt();
   };