diff --git a/python/paddle/fluid/tests/unittests/test_unstack_op.py b/python/paddle/fluid/tests/unittests/test_unstack_op.py index 1dda05fb0a6b8669577c70134c4acb73afb51621..745e14983a56a4c8456b9fe4897378f42543ab94 100755 --- a/python/paddle/fluid/tests/unittests/test_unstack_op.py +++ b/python/paddle/fluid/tests/unittests/test_unstack_op.py @@ -84,5 +84,29 @@ class TestStackOp6(TestUnStackOpBase): self.axis = 2 +class TestUnstackZeroInputOp(unittest.TestCase): + def unstack_zero_input_static(self): + + paddle.enable_static() + + array = np.array([], dtype=np.float32) + x = paddle.to_tensor(np.reshape(array, [0]), dtype='float32') + paddle.unstack(x, axis=1) + + def unstack_zero_input_dynamic(self): + + array = np.array([], dtype=np.float32) + x = paddle.to_tensor(np.reshape(array, [0]), dtype='float32') + paddle.unstack(x, axis=1) + + def test_type_error(self): + paddle.disable_static() + + self.assertRaises(ValueError, self.unstack_zero_input_dynamic) + self.assertRaises(ValueError, self.unstack_zero_input_static) + + paddle.disable_static() + + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py index bdd903ee8f1966179301bccab04e3249b09ae299..923e6923d6d63c66af55cb7f360411ce8bf88684 100644 --- a/python/paddle/tensor/manipulation.py +++ b/python/paddle/tensor/manipulation.py @@ -543,6 +543,10 @@ def unstack(x, axis=0, num=None): y = paddle.unstack(x, axis=1) # unstack with second axis, which results 3 tensors with shape=[2, 5] """ + if not (-x.ndim <= axis < x.ndim): + raise ValueError( + '`axis` must be in the range [-{0}, {0})'.format(x.ndim) + ) if in_dygraph_mode(): if num is None: num = x.shape[axis]