add dtype for unique (#26655)

* update doc, test=document_fix

* add attr(dtype)

* refine code

paddle/fluid/operators/unique_op.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/unique_op.h"
+#include "paddle/fluid/framework/op_version_registry.h"
 
 namespace paddle {
 namespace operators {
@@ -149,3 +150,34 @@ REGISTER_OP_CPU_KERNEL(
     ops::UniqueKernel<paddle::platform::CPUDeviceContext, double>,
     ops::UniqueKernel<paddle::platform::CPUDeviceContext, int32_t>,
     ops::UniqueKernel<paddle::platform::CPUDeviceContext, int64_t>);
+REGISTER_OP_VERSION(unique)
+    .AddCheckpoint(
+        R"ROC(
+        Upgrade unique, add 2 outputs [Indices, Counts] and 5 attribute
+        [return_index, return_inverse, return_counts, axis, is_sorted].
+      )ROC",
+        paddle::framework::compatible::OpVersionDesc()
+            .NewOutput("Indices",
+                       "The indices of the input tensor that result in the "
+                       "unique tensor.")
+            .NewOutput("Counts", "The counts for each unique element.")
+            .NewAttr("return_index",
+                     "If True, also return the indices of the input"
+                     " tensor that result in the unique Tensor.",
+                     false)
+            .NewAttr("return_inverse",
+                     "If True, also return the indices for where elements"
+                     " in the original input ended up in the returned unique "
+                     "tensor.",
+                     false)
+            .NewAttr("return_counts",
+                     "If True, also return the counts for each unique element.",
+                     false)
+            .NewAttr("axis",
+                     "The axis to apply unique. If None, the input will be "
+                     "flattened.",
+                     {})
+            .NewAttr("is_sorted",
+                     "If True, the unique elements of X are in ascending order."
+                     "Otherwise, the unique elements are not sorted.",
+                     false));

python/paddle/fluid/layers/nn.py

@@ -14098,17 +14098,11 @@ def sign(x):
 
 def unique(x, dtype='int32'):
     """
-    :alias_main: paddle.unique
-	:alias: paddle.unique,paddle.tensor.unique,paddle.tensor.manipulation.unique
-	:old_api: paddle.fluid.layers.unique
-
-    **unique**
-
     Return a unique tensor for `x` and an index tensor pointing to this unique tensor.
 
     Args:
-        x(Variable): A 1-D input tensor.
-        dtype(np.dtype|core.VarDesc.VarType|str): The type of index tensor: int32, int64.
+        x(Tensor): A 1-D input tensor, it's data type should be float32, float64, int32, int64.
+        dtype(np.dtype|str, optional): The type of index tensor: int32, int64. Default: int32.
 
     Returns:
         tuple: (out, index). `out` is the unique tensor for `x`, with identical dtype to `x`, and \

python/paddle/fluid/tests/unittests/test_unique.py

@@ -233,6 +233,24 @@ class TestUniqueAPI(unittest.TestCase):
         self.assertTrue((counts.numpy() == np_counts).all(), True)
         paddle.enable_static()
 
+    def test_dygraph_attr_dtype(self):
+        paddle.disable_static()
+        x_data = x_data = np.random.randint(0, 10, (120))
+        x = paddle.to_tensor(x_data)
+        out, indices, inverse, counts = paddle.unique(
+            x,
+            return_index=True,
+            return_inverse=True,
+            return_counts=True,
+            dtype="int32")
+        expected_out, np_indices, np_inverse, np_counts = np.unique(
+            x_data, return_index=True, return_inverse=True, return_counts=True)
+        self.assertTrue((out.numpy() == expected_out).all(), True)
+        self.assertTrue((indices.numpy() == np_indices).all(), True)
+        self.assertTrue((inverse.numpy() == np_inverse).all(), True)
+        self.assertTrue((counts.numpy() == np_counts).all(), True)
+        paddle.enable_static()
+
     def test_static_graph(self):
         with paddle.static.program_guard(paddle.static.Program(),
                                          paddle.static.Program()):
@@ -282,6 +300,9 @@ class TestUniqueError(unittest.TestCase):
         def test_axis():
             result = paddle.unique(x, axis='12')
 
+        def test_dtype():
+            result = paddle.unique(x, dtype='float64')
+
         self.assertRaises(TypeError, test_axis)
 
 

python/paddle/tensor/manipulation.py

@@ -612,6 +612,7 @@ def unique(x,
            return_inverse=False,
            return_counts=False,
            axis=None,
+           dtype="int64",
            name=None):
     """
     Returns the unique elements of `x` in ascending order.
@@ -625,6 +626,8 @@ def unique(x,
         return_counts(bool, optional): If True, also return the counts for each unique element.
         axis(int, optional): The axis to apply unique. If None, the input will be flattened.
             Default: None.
+        dtype(np.dtype|str, optional): The date type of `indices` or `inverse` tensor: int32 or int64.
+            Default: int64.
         name(str, optional): Name for the operation. For more information, please refer to
             :ref:`api_guide_Name`. Default: None.
 
@@ -650,6 +653,7 @@ def unique(x,
             np_counts = counts.numpy() # [1 1 3 1]
 
             x_data = np.array([[2, 1, 3], [3, 0, 1], [2, 1, 3]])
+            x = paddle.to_tensor(x_data)
             unique = paddle.unique(x)
             np_unique = unique.numpy() # [0 1 2 3]
 
@@ -662,11 +666,10 @@ def unique(x,
         axis = []
     else:
         axis = [axis]
-
+    attr_dtype = convert_np_dtype_to_dtype_(dtype)
     if in_dygraph_mode():
         out, inverse, indices, counts = core.ops.unique(
-            x, 'dtype',
-            convert_np_dtype_to_dtype_('int32'), 'return_index', return_index,
+            x, 'dtype', attr_dtype, 'return_index', return_index,
             'return_inverse', return_inverse, 'return_counts', return_counts,
             'axis', axis, "is_sorted", True)
         outs = [out]
@@ -687,12 +690,13 @@ def unique(x,
     check_type(return_index, 'return_index', bool, 'unique')
     check_type(return_inverse, 'return_inverse', bool, 'unique')
     check_type(return_counts, 'return_counts', bool, 'unique')
+    check_dtype(dtype, 'dtype', ['int32', 'int64'], 'unique')
     if len(axis) != 0:
         check_type(axis[0], 'axis', int, 'unique')
 
     helper = LayerHelper('unique', **locals())
     attrs = {
-        'dtype': int(core.VarDesc.VarType.INT32),
+        'dtype': attr_dtype,
         "return_index": return_index,
         "return_inverse": return_inverse,
         "return_counts": return_counts,
@@ -702,19 +706,19 @@ def unique(x,
     out = helper.create_variable_for_type_inference(
         dtype=x.dtype, stop_gradient=True)
     inverse = helper.create_variable_for_type_inference(
-        dtype=core.VarDesc.VarType.INT64, stop_gradient=True)
+        dtype=attr_dtype, stop_gradient=True)
     outputs = {"Out": out, "Index": inverse}
     outs = [out]
     if return_index:
         indices = helper.create_variable_for_type_inference(
-            dtype=core.VarDesc.VarType.INT64, stop_gradient=True)
+            dtype=attr_dtype, stop_gradient=True)
         outputs["Indices"] = indices
         outs.append(indices)
     if return_inverse:
         outs.append(inverse)
     if return_counts:
         counts = helper.create_variable_for_type_inference(
-            dtype=core.VarDesc.VarType.INT64, stop_gradient=True)
+            dtype=attr_dtype, stop_gradient=True)
         outputs["Counts"] = counts
         outs.append(counts)