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

python/paddle/__init__.py

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

python/paddle/nn/functional/__init__.py

@@ -78,7 +78,8 @@
     conv3d,
     conv3d_transpose,
 )
-from .distance import pairwise_distance
+
+from .distance import pairwise_distance, pdist
 from .extension import (
     diag_embed,  # noqa: F401
     gather_tree,
@@ -162,6 +163,7 @@
     'conv3d',
     'conv3d_transpose',
     'pairwise_distance',
+    'pdist',
     'elu',
     'elu_',
     'gelu',

python/paddle/nn/functional/distance.py

@@ -106,3 +106,40 @@ def pairwise_distance(x, y, p=2.0, epsilon=1e-6, keepdim=False, name=None):
         )
 
         return out
+
+
+def pdist(x, p=2.0, 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`.
+        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)
+            >>> print(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.linalg.norm(x[..., None, :] - x[..., None, :, :], p=p, axis=-1)
+    mask = ~paddle.tril(paddle.ones(d.shape, dtype='bool'))
+    return paddle.masked_select(d, mask)

test/legacy_test/test_pdist.py

@@ -0,0 +1,151 @@
+#   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.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,
+            )
+            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,
+        )
+        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')
+
+
+class TestPdistAPICase8(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,
+            )
+            exe = paddle.static.Executor(self.place)
+            res = exe.run(feed={'x': self.x}, fetch_list=[out0])
+            out_ref = ref_pdist(self.x, self.p)
+            np.testing.assert_allclose(out_ref, res[0])
+
+    def test_dygraph_api(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x)
+        out0 = paddle.pdist(
+            x,
+            self.p,
+        )
+        out_ref = ref_pdist(self.x, self.p)
+        np.testing.assert_allclose(out_ref, out0.numpy())
+        paddle.enable_static()
+
+
+class TestPdistShapeError(unittest.TestCase):
+    def test_error(self):
+        with self.assertRaises(AssertionError):
+            self.x = np.random.rand(50, 10, 20).astype('float64')
+            self.p = 2.0
+            x = paddle.to_tensor(self.x)
+            out0 = paddle.pdist(
+                x,
+                self.p,
+            )
+
+
+if __name__ == '__main__':
+    paddle.enable_static()
+    unittest.main()