Merge pull request #49843 from geetachavan1/cherrypicks_22RRN

Fix `tf.raw_ops.RaggedTensorToVariant` invalid resize.
tensorflow · May 31, 2021 · 4e8ff09 · 4e8ff09
2 parents e5d259f + 1650bb0
commit 4e8ff09
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
@@ -175,6 +175,11 @@ class RaggedTensorToVariantOp : public OpKernel {
 
     // Unbatch the Ragged Tensor and encode the components.
     std::vector<RaggedTensor> ragged_components;
+    auto batched_splits_top_vec =
+        batched_ragged_input.splits(0).vec<SPLIT_TYPE>();
+    int num_components = batched_splits_top_vec.size() - 1;
+    OP_REQUIRES(context, num_components >= 0,
+                errors::Internal("Invalid split argument."));
     OP_REQUIRES_OK(context, UnbatchRaggedZerothDim<VALUE_TYPE, SPLIT_TYPE>(
                                 batched_ragged_input, &ragged_components));
     std::vector<Tensor> encoded_components(ragged_components.size());