diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc index 8a3c33a4d6c72..3b3202c291725 100644 --- a/paddle/phi/infermeta/unary.cc +++ b/paddle/phi/infermeta/unary.cc @@ -4648,6 +4648,7 @@ void UniqueRawInferMeta(const MetaTensor& x, if (axis_value < 0) { axis_value += x.dims().size(); } + PADDLE_ENFORCE_LT( axis_value, x.dims().size(), @@ -4655,6 +4656,14 @@ void UniqueRawInferMeta(const MetaTensor& x, "the dimension size(%d) of x.", axis_value, x.dims().size())); + PADDLE_ENFORCE_GE( + axis_value, + 0, + phi::errors::InvalidArgument( + "The axis(%d) + rank(x) (%d) should be greater than or equal to 0.", + axis_value, + -x.dims().size())); + auto out_dims = x.dims(); out_dims[axis_value] = -1; out->set_dims(out_dims); diff --git a/python/paddle/fluid/tests/unittests/test_unique.py b/python/paddle/fluid/tests/unittests/test_unique.py index 9183c1bd676bb..b3ae10a6c335e 100644 --- a/python/paddle/fluid/tests/unittests/test_unique.py +++ b/python/paddle/fluid/tests/unittests/test_unique.py @@ -190,6 +190,32 @@ def init_config(self): } +class TestUniqueOpAxisNeg(TestUniqueOp): + def init_config(self): + self.inputs = {'X': np.random.random((6, 1, 8)).astype('float64')} + unique, indices, inverse, counts = np.unique( + self.inputs['X'], + return_index=True, + return_inverse=True, + return_counts=True, + axis=-1, + ) + self.attrs = { + 'dtype': int(core.VarDesc.VarType.INT32), + "return_index": True, + "return_inverse": True, + "return_counts": True, + "axis": [-1], + "is_sorted": True, + } + self.outputs = { + 'Out': unique, + 'Indices': indices, + "Index": inverse, + "Counts": counts, + } + + class TestUniqueOpAxis1(TestUniqueOp): def init_config(self): self.inputs = {'X': np.random.random((3, 8, 8)).astype('float64')}