Merge pull request #53695 from yongtang:53660-tf.sparse.split-crash

tensorflower-gardener · mihaimaruseac · commit 894ba6b75bc8 · 2022-01-25T08:40:50.000-08:00
PiperOrigin-RevId: 420811652
Change-Id: I83742482770ba0bf7c3ccd57508c40fb9cdbe2f7
diff --git a/tensorflow/core/kernels/sparse_split_op.cc b/tensorflow/core/kernels/sparse_split_op.cc
@@ -30,11 +30,16 @@ class SparseSplitOp : public OpKernel {
   }
 
   void Compute(OpKernelContext* context) override {
-    const int64_t axis_input = context->input(0).scalar<int64_t>()();
+    const Tensor& input_axis = context->input(0);
     const Tensor& input_indices = context->input(1);
     const Tensor& input_values = context->input(2);
     const Tensor& input_shape = context->input(3);
 
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(input_axis.shape()),
+                    errors::InvalidArgument(
+                        "Input axis should be a scalar but received shape ",
+                        input_axis.shape().DebugString()),
+                    done);
     OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices.shape()),
                 errors::InvalidArgument(
                     "Input indices should be a matrix but received shape ",
@@ -48,6 +53,7 @@ class SparseSplitOp : public OpKernel {
                     "Input shape should be a vector but received shape ",
                     input_shape.shape().DebugString()));
 
+    const int64_t axis_input = input_axis.scalar<int64_t>()();
     const int64_t input_rank = input_shape.vec<int64_t>().size();
     const int64_t axis =
         (axis_input < 0) ? input_rank + axis_input : axis_input;
diff --git a/tensorflow/python/kernel_tests/sparse_split_op_test.py b/tensorflow/python/kernel_tests/sparse_split_op_test.py
@@ -257,6 +257,15 @@ def testArgumentErrors(self):
     with self.assertRaisesRegex(ValueError, 'axis is required'):
       sparse_ops.sparse_split(num_split=2, sp_input=1)
 
+  def testInvalidArgumentError(self):
+    # Test case for GitHub issue 53660.
+    axis = [1, 2]
+    with self.assertRaisesRegexp(errors.InvalidArgumentError,
+                                 r'axis should be a scalar'):
+      self.evaluate(
+          sparse_ops.sparse_split(
+              sp_input=self._SparseTensor_4x6(), num_split=3, axis=axis))
+
 
 if __name__ == '__main__':
   test.main()
