System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Y
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.7.0
• Python version: 3.8
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source):N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

Standalone code to reproduce the issue

import tensorflow as tf
data = tf.random.uniform([1, 32, 32], dtype=tf.float32)
axis = [1, 2]
x = tf.sparse.from_dense(data)
result = tf.sparse.split(x,3, axis=axis) # crash

Session crashes. tf.sparse.split failed to do proper checking for axis.

Describe the expected behavior
tf.sparse.split should raise InvalidArgumentError when axis is not a 0-D tensor, instead of crashing.