【Hackathon 5th No.17】 为 Paddle 新增 pdist API -part

python/paddle/__init__.py

@@ -524,6 +524,10 @@
     flops,
 )
 
+from .nn.functional.distance import (  # noqa: F401
+    pdist,
+)
+
 import paddle.text  # noqa: F401
 import paddle.vision  # noqa: F401
 
@@ -702,6 +706,7 @@
     'sin_',
     'dist',
     'cdist',
+    'pdist',
     'unbind',
     'meshgrid',
     'arange',

python/paddle/nn/functional/__init__.py

@@ -78,7 +78,7 @@
     conv3d,
     conv3d_transpose,
 )
-from .distance import pairwise_distance
+from .distance import pairwise_distance, pdist
 from .extension import diag_embed, gather_tree, sequence_mask, temporal_shift
 from .flash_attention import (  # noqa: F401
     scaled_dot_product_attention,
@@ -157,6 +157,7 @@
     'conv3d',
     'conv3d_transpose',
     'pairwise_distance',
+    'pdist',
     'elu',
     'elu_',
     'gelu',

python/paddle/nn/functional/distance.py

@@ -106,3 +106,49 @@ def pairwise_distance(x, y, p=2.0, epsilon=1e-6, keepdim=False, name=None):
         )
 
         return out
+
+
+def pdist(
+    x, p=2.0, compute_mode="use_mm_for_euclid_dist_if_necessary", name=None
+):
+    r'''
+    Computes the p-norm distance between every pair of row vectors in the input.
+
+    Args:
+        x (Tensor): A tensor with shape :math:`N \times M`.
+        p (float, optional): The value for the p-norm distance to calculate between each vector pair. Default: :math:`2.0`.
+        compute_mode (str, optional): The mode for compute distance.
+
+            - ``use_mm_for_euclid_dist_if_necessary`` , for p = 2.0 and (P > 25 or R > 25), it will use matrix multiplication to calculate euclid distance if possible.
+            - ``use_mm_for_euclid_dist`` , for p = 2.0, it will use matrix multiplication to calculate euclid distance.
+            - ``donot_use_mm_for_euclid_dist`` , it will not use matrix multiplication to calculate euclid distance.
+
+            Default: ``use_mm_for_euclid_dist_if_necessary``.
+        name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None.
+
+    Returns:
+        Tensor with shape: math:`N(N-1)/2` the dtype is same as input tensor.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> a = paddle.randn([4, 5])
+            >>> a
+            Tensor(shape=[4, 5], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+                   [[-0.33173719, -0.93648648, -0.01741328, -0.94435263,  2.22178721],
+                    [-0.65466857,  0.10307083,  0.08741203, -0.91078597,  0.72589827],
+                    [ 0.06907391, -0.27584535,  1.35355449, -0.69688839,  0.18408430],
+                    [-0.00939178, -0.32901841, -1.06503606,  0.81856263,  0.16791444]])
+            >>> pdist_out=paddle.pdist(a)
+            >>> pdist_out
+            Tensor(shape=[6], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+                   [1.85331142, 2.58652687, 2.98273396, 1.61549115, 2.28762150, 2.85576940])
+
+    '''
+
+    x_shape = list(x.shape)
+    assert len(x_shape) == 2, "The x must be 2-dimensional"
+    d = paddle.cdist(x, x, p, compute_mode)
+    mask = ~paddle.tril(paddle.ones(d.shape, dtype='bool'))
+    return paddle.masked_select(d, mask)

test/legacy_test/test_pdist.py

@@ -0,0 +1,143 @@
+#   Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+def ref_pdist(x, p=2.0):
+    dist = np.linalg.norm(x[..., None, :] - x[None, :, :], ord=p, axis=-1)
+    res = []
+    rows, cols = dist.shape
+    for i in range(rows):
+        for j in range(cols):
+            if i >= j:
+                continue
+            res.append(dist[i][j])
+    return np.array(res)
+
+
+class TestpdistAPI(unittest.TestCase):
+    def setUp(self):
+        self.x = np.random.rand(10, 20).astype('float32')
+        self.p = 2.0
+        self.compute_mode = "use_mm_for_euclid_dist_if_necessary"
+        self.init_input()
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+    def init_input(self):
+        pass
+
+    def test_static_api(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program()):
+            x = paddle.static.data('x', self.x.shape, dtype=self.x.dtype)
+            out = paddle.pdist(x, self.p, self.compute_mode)
+            exe = paddle.static.Executor(self.place)
+            res = exe.run(feed={'x': self.x}, fetch_list=[out])
+            out_ref = ref_pdist(self.x, self.p)
+            np.testing.assert_allclose(out_ref, res[0], rtol=1e-5, atol=1e-5)
+
+    def test_dygraph_api(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x)
+        out = paddle.pdist(x, self.p, self.compute_mode)
+        out_ref = ref_pdist(self.x, self.p)
+        np.testing.assert_allclose(out_ref, out.numpy(), rtol=1e-5, atol=1e-5)
+        paddle.enable_static()
+
+
+class TestpdistAPICase1(TestpdistAPI):
+    def init_input(self):
+        self.p = 0
+
+
+class TestpdistAPICase2(TestpdistAPI):
+    def init_input(self):
+        self.p = 1.0
+
+
+class TestpdistAPICase3(TestpdistAPI):
+    def init_input(self):
+        self.p = 3.0
+
+
+class TestpdistAPICase4(TestpdistAPI):
+    def init_input(self):
+        self.p = 1.5
+
+
+class TestpdistAPICase5(TestpdistAPI):
+    def init_input(self):
+        self.p = 2.5
+
+
+class TestpdistAPICase6(TestpdistAPI):
+    def init_input(self):
+        self.p = float('inf')
+
+
+class TestpdistAPICase7(TestpdistAPI):
+    def init_input(self):
+        self.x = np.random.rand(50, 20).astype('float64')
+        self.compute_mode = "use_mm_for_euclid_dist"
+
+
+class TestpdistAPICase8(TestpdistAPI):
+    def init_input(self):
+        self.x = np.random.rand(50, 20).astype('float64')
+        self.compute_mode = "donot_use_mm_for_euclid_dist"
+
+
+class TestpdistAPICase9(TestpdistAPI):
+    def init_input(self):
+        self.x = np.random.rand(500, 100).astype('float64')
+
+    def test_static_api(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program()):
+            x = paddle.static.data('x', self.x.shape, dtype=self.x.dtype)
+            out0 = paddle.pdist(x, self.p, self.compute_mode)
+            out1 = paddle.pdist(x, self.p, "donot_use_mm_for_euclid_dist")
+            out2 = paddle.pdist(x, self.p, "use_mm_for_euclid_dist")
+            exe = paddle.static.Executor(self.place)
+            res = exe.run(feed={'x': self.x}, fetch_list=[out0, out1, out2])
+            out_ref = ref_pdist(self.x, self.p)
+            np.testing.assert_allclose(out_ref, res[0])
+            np.testing.assert_allclose(out_ref, res[1])
+            np.testing.assert_allclose(out_ref, res[2])
+
+    def test_dygraph_api(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x)
+        out0 = paddle.pdist(x, self.p, self.compute_mode)
+        out1 = paddle.pdist(x, self.p, "donot_use_mm_for_euclid_dist")
+        out2 = paddle.pdist(x, self.p, "use_mm_for_euclid_dist")
+        out_ref = ref_pdist(self.x, self.p)
+        np.testing.assert_allclose(out_ref, out0.numpy())
+        np.testing.assert_allclose(out_ref, out1.numpy())
+        np.testing.assert_allclose(out_ref, out2.numpy())
+        paddle.enable_static()
+
+
+if __name__ == '__main__':
+    paddle.enable_static()
+    unittest.main()