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

tensorflower-gardener · tensorflower-gardener · commit 61bf91e76817 · 2022-01-10T11:35:05.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
@@ -74,11 +74,16 @@ void SparseSplitOpImpl(OpKernelContext* context, int num_split,
     done = [] {};
   }
 
-  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_ASYNC(context, TensorShapeUtils::IsScalar(input_axis.shape()),
+                    errors::InvalidArgument(
+                        "Input axis should be a scalar but received shape ",
+                        input_axis.shape().DebugString()),
+                    done);
   OP_REQUIRES_ASYNC(context, TensorShapeUtils::IsMatrix(input_indices.shape()),
                     errors::InvalidArgument(
                         "Input indices should be a matrix but received shape ",
@@ -95,6 +100,7 @@ void SparseSplitOpImpl(OpKernelContext* context, int num_split,
                         input_shape.shape().DebugString()),
                     done);
 
+  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_ops/sparse_split_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_split_op_test.py
@@ -279,6 +279,15 @@ def testSplitEmpty(self):
     self.assertAllEqual(sparse_splits1[1].values, [])
     self.assertAllEqual(sparse_splits1[1].dense_shape, [4, 3])
 
+  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()
