Fix checkfail in UnicodeEncode Op

The Op UnicodeEncode segfaults when passed 2D tensor to `input_splits`. It has the below check in SetShapeFn which supposed to raise exception if rank !=1 AFAIk. This seems not working for reason unknown to me. https://github.com/tensorflow/tensorflow/blob/6f64ad5d767a034df45a5eaab8b36fd688cd1217/tensorflow/core/ops/string_ops.cc#L316-L317 Same with input_values argument also. Added an explicit check in Op.
tensorflow · Mar 11, 2024 · 176048a · 176048a
1 parent 6f64ad5
commit 176048a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/unicode_ops.cc b/tensorflow/core/kernels/unicode_ops.cc
@@ -532,6 +532,10 @@ class UnicodeEncodeOp : public OpKernel {
     const Tensor& input_splits = context->input(1);
     const auto input_splits_flat = input_splits.flat<SPLITS_TYPE>();
 
+    OP_REQUIRES(
+        context, input_tensor.dims() == 1 && input_splits.dims() == 1,
+        absl::InvalidArgumentError(
+            "Both the input_tensor and input_splits should be of rank 1. "));   
     OP_REQUIRES(
         context, input_splits.NumElements() > 0,
         errors::InvalidArgument("Input_splits should contain elements, but "