Add value check for chunk

python/paddle/tensor/manipulation.py

@@ -4482,6 +4482,13 @@ def chunk(
  >>> # out2.shape [3, 3, 5]
  """
  check_type(chunks, 'chunks', (int), 'chunk')
+ # check the chunks value to avoid the meanless split operation, such as <=0 or > x.shape[axis]
+ if chunks <= 0 or (
+ isinstance(axis, int) and axis >= 0 and chunks > x.shape[axis]
+ ):
+ raise ValueError(
+ f"The value of 'chunks' must be greater than 0 and less than or equal to the size of the dimension specified by 'axis', but received chunks={chunks} and x.shape[axis]={x.shape[axis]}."
+ )
  return split(x, num_or_sections=chunks, axis=axis, name=name)
 
 

test/legacy_test/test_chunk_op.py

@@ -60,6 +60,19 @@ def test_0_chunks_tensor():
 
  self.assertRaises(ValueError, test_0_chunks_tensor)
 
+ def test_negative_chunks_tensor():
+ x = paddle.uniform([2, 3, 4], dtype='float32')
+ paddle.chunk(x, chunks=-1)
+
+ self.assertRaises(ValueError, test_negative_chunks_tensor)
+
+ def test_chunks_greater_than_dim():
+ x = paddle.uniform([2, 3, 4], dtype='float32')
+ # axis=1, shape=3, chunks=5 > 3
+ paddle.chunk(x, chunks=5, axis=1)
+
+ self.assertRaises(ValueError, test_chunks_greater_than_dim)
+
 
 class API_TestChunk(unittest.TestCase):
  def test_out(self):