【Hackathon 5th No.115】为 paddle.linalg.norm 进行功能对齐与功能增强-增加nuclear_norm (#60766)

Author: zbt78

python/paddle/tensor/linalg.py

@@ -295,8 +295,8 @@ def norm(x, p='fro', axis=None, keepdim=False, name=None):
     Args:
         x (Tensor): The input tensor could be N-D tensor, and the input data
             type could be float32 or float64.
-        p (float|string, optional): Order of the norm. Supported values are `fro`, `0`, `1`, `2`,
-            `inf`, `-inf` and any positive real number yielding the corresponding p-norm. Not supported: ord < 0 and nuclear norm.
+        p (float|string, optional): Order of the norm. Supported values are `fro`, `nuc`, `0`, `1`, `2`,
+            `inf`, `-inf` and any positive real number yielding the corresponding p-norm. Not supported: ord < 0.
             Default value is `fro`.
         axis (int|list|tuple, optional): The axis on which to apply norm operation. If axis is int
             or list(int)/tuple(int)  with only one element, the vector norm is computed over the axis.
@@ -374,6 +374,21 @@ def norm(x, p='fro', axis=None, keepdim=False, name=None):
              [4., 3., 2., 1.]])
     """
 
+    def _backshift_permutation(dim0, dim1, dimn):
+        """
+        Auxiliary function for matrix_norm
+        Computes the permutation that moves the two given dimensions to the back
+        """
+        ret = [i for i in range(dimn) if i != dim0 and i != dim1]
+        ret.extend((dim0, dim1))
+        return ret
+
+    def _inverse_permutation(perm):
+        """
+        Given a permutation, returns its inverse. It's equivalent to argsort on an array
+        """
+        return [i for i, j in sorted(enumerate(perm), key=lambda ij: ij[1])]
+
     def frobenius_norm(input, dim=None, keepdim=False, name=None):
         """
         The frobenius norm OP is to calculate the frobenius norm of certain two dimensions of Tensor `input`.
@@ -414,6 +429,98 @@ def frobenius_norm(input, dim=None, keepdim=False, name=None):
             )
             return out
 
+    def nuclear_norm(input, axis=axis, keepdim=False, name=None):
+        """
+        The nuclear norm OP is to calculate the nuclear norm of certain two dimensions of Tensor `input`.
+        Args:
+          input (Variable): Tensor, data type float32, float64.
+          dim (list): Two dimensions.
+          keepdim (bool, optional): Whether keep the dimensions as the `input`, Default False.
+          name (str, optional): The default value is None. Normally there is no need for
+              user to set this property. For more information, please refer to :ref:`api_guide_Name`.
+        """
+
+        perm = _backshift_permutation(axis[0], axis[1], len(input.shape))
+        inv_perm = _inverse_permutation(perm)
+
+        if in_dynamic_mode():
+            transposed = _C_ops.transpose(input, perm)
+            u, s, vh = _C_ops.svd(transposed, False)
+            result = _C_ops.sum(s, -1, None, keepdim)
+            if keepdim:
+                result = _C_ops.transpose(
+                    _C_ops.unsqueeze(result, -1), inv_perm
+                )
+            return result
+
+        attrs = {'axis': axis, 'keepdim': keepdim}
+
+        check_variable_and_dtype(
+            input, 'input', ['float32', 'float64'], 'nuclear_norm'
+        )
+
+        block = LayerHelper('nuclear_nrom', **locals())
+        out = block.create_variable_for_type_inference(
+            dtype=block.input_dtype()
+        )
+
+        transpose_out = block.create_variable_for_type_inference(
+            dtype=block.input_dtype()
+        )
+        input_shape = block.create_variable_for_type_inference(
+            dtype=block.input_dtype()
+        )
+
+        block.append_op(
+            type='transpose2',
+            inputs={'X': [input]},
+            outputs={'Out': [transpose_out], 'XShape': [input_shape]},
+            attrs={'axis': perm},
+        )
+
+        u = block.create_variable_for_type_inference(dtype=block.input_dtype())
+        s = block.create_variable_for_type_inference(dtype=block.input_dtype())
+        vt = block.create_variable_for_type_inference(dtype=block.input_dtype())
+        block.append_op(
+            type='svd',
+            inputs={'X': [transpose_out]},
+            outputs={'U': u, 'VH': vt, 'S': s},
+            attrs={'full_matrices': False},
+        )
+
+        reduce_all, sum_axis = _get_reduce_axis(-1, s)
+        block.append_op(
+            type='reduce_sum',
+            inputs={'X': s},
+            outputs={'Out': out},
+            attrs={
+                'dim': sum_axis,
+                'keep_dim': keepdim,
+                'reduce_all': reduce_all,
+            },
+        )
+
+        if keepdim:
+            unsqueeze_out = block.create_variable_for_type_inference(
+                dtype=block.input_dtype()
+            )
+
+            block.append_op(
+                type='unsqueeze2',
+                inputs={'X': [out]},
+                outputs={'Out': [unsqueeze_out], 'XShape': [input_shape]},
+                attrs={'axes': [-1]},
+            )
+
+            block.append_op(
+                type='transpose2',
+                inputs={'X': [unsqueeze_out]},
+                outputs={'Out': [out], 'XShape': [input_shape]},
+                attrs={'axis': inv_perm},
+            )
+
+        return out
+
     def vector_norm(
         input, porder=None, axis=None, keepdim=False, asvector=False, name=None
     ):
@@ -616,6 +723,8 @@ def p_matrix_norm(input, porder=1.0, axis=axis, keepdim=False, name=None):
     elif isinstance(axis, list) and len(axis) == 2:
         if p == "fro":
             return frobenius_norm(x, dim=axis, keepdim=keepdim, name=name)
+        elif p == "nuc":
+            return nuclear_norm(x, axis=axis, keepdim=keepdim, name=name)
         elif p == np.inf or p == -np.inf:
             return inf_norm(x, porder=p, axis=axis, keepdim=keepdim, name=name)
         elif p == 0:

test/legacy_test/test_norm_all.py

@@ -80,10 +80,23 @@ def numpy_frobenius_norm(x, axis=None, keepdims=False):
     return r
 
 
+def numpy_nuclear_norm(x, axis=None, keepdims=False):
+    if isinstance(axis, list):
+        axis = tuple(axis)
+    r = np.linalg.norm(x, ord='nuc', axis=axis, keepdims=keepdims).astype(
+        x.dtype
+    )
+    return r
+
+
 def frobenius_norm(x, dim, keep_dim, reduce_all):
     return paddle.linalg.norm(x, p='fro', axis=dim, keepdim=keep_dim)
 
 
+def nuclear_norm(x, dim, keep_dim, reduce_all):
+    return paddle.linalg.norm(x, p='nuc', axis=dim, keepdim=keep_dim)
+
+
 class TestFrobeniusNormOp(OpTest):
     def setUp(self):
         self.python_api = frobenius_norm
@@ -433,6 +446,33 @@ def run_fro(self, p, axis, shape_x, dtype, keep_dim, check_dim=False):
         )
 
 
+def check_nuc_static(self, p, axis, shape_x, dtype, keep_dim, check_dim=False):
+    with base.program_guard(base.Program()):
+        data = paddle.static.data(name="X", shape=shape_x, dtype=dtype)
+        out = paddle.norm(x=data, p=p, axis=axis, keepdim=keep_dim)
+        place = base.CPUPlace()
+        exe = base.Executor(place)
+        np_input = (np.random.rand(*shape_x) + 1.0).astype(dtype)
+        expected_result = numpy_nuclear_norm(
+            np_input, axis=axis, keepdims=keep_dim
+        )
+        (result,) = exe.run(feed={"X": np_input}, fetch_list=[out])
+    np.testing.assert_allclose(result, expected_result, rtol=1e-6, atol=1e-8)
+    if keep_dim and check_dim:
+        np.testing.assert_equal(result.shape, expected_result.shape)
+
+
+def check_nuc_dygraph(self, p, axis, shape_x, dtype, keep_dim, check_dim=False):
+    x_numpy = (np.random.random(shape_x) + 1.0).astype(dtype)
+    expected_result = numpy_nuclear_norm(x_numpy, axis, keep_dim)
+    x_paddle = paddle.to_tensor(x_numpy)
+    result = paddle.norm(x=x_paddle, p=p, axis=axis, keepdim=keep_dim)
+    result = result.numpy()
+    np.testing.assert_allclose(result, expected_result, rtol=1e-6, atol=1e-8)
+    if keep_dim and check_dim:
+        np.testing.assert_equal(result.shape, expected_result.shape)
+
+
 def run_pnorm(self, p, axis, shape_x, dtype, keep_dim, check_dim=False):
     with base.program_guard(base.Program()):
         data = paddle.static.data(name="X", shape=shape_x, dtype=dtype)
@@ -469,6 +509,8 @@ def run_graph(self, p, axis, shape_x, dtype):
     out_fro = paddle.norm(x, p='fro')
     out_fro = paddle.norm(x, p='fro', axis=0)
     out_fro = paddle.norm(x, p='fro', axis=[0, 1])
+    # compute nuclear norm.
+    out_nuc = paddle.norm(x, p='nuc', axis=[0, 1])
     # compute 2-order  norm along [0,1] dimension.
     out_pnorm = paddle.norm(x, p=2, axis=[0, 1])
     out_pnorm = paddle.norm(x, p=2)
@@ -508,6 +550,15 @@ def test_basic(self):
                 keep_dim=keep,
                 check_dim=True,
             )
+            check_nuc_static(
+                self,
+                p='nuc',
+                axis=[0, 1],
+                shape_x=[2, 3, 4],
+                dtype='float64',
+                keep_dim=keep,
+                check_dim=True,
+            )
             run_pnorm(
                 self,
                 p=2,
@@ -636,13 +687,38 @@ def test_basic(self):
     def test_dygraph(self):
         run_graph(self, p='fro', axis=None, shape_x=[2, 3, 4], dtype="float32")
 
+        paddle.disable_static()
+        keep_dims = {False, True}
+        for keep in keep_dims:
+            check_nuc_dygraph(
+                self,
+                p='nuc',
+                axis=[0, 1],
+                shape_x=[2, 3, 4],
+                dtype='float64',
+                keep_dim=keep,
+                check_dim=True,
+            )
+            check_nuc_dygraph(
+                self,
+                p='nuc',
+                axis=[1, 2],
+                shape_x=[2, 3, 4, 5],
+                dtype='float64',
+                keep_dim=keep,
+                check_dim=True,
+            )
+        paddle.enable_static()
+
     def test_name(self):
         with base.program_guard(base.Program()):
             x = paddle.static.data(name="x", shape=[10, 10], dtype="float32")
             y_1 = paddle.norm(x, p='fro', name='frobenius_name')
             y_2 = paddle.norm(x, p=2, name='pnorm_name')
+            y_3 = paddle.norm(x, p='nuc', axis=[0, 1], name='nuclear_name')
             self.assertEqual(('frobenius_name' in y_1.name), True)
             self.assertEqual(('pnorm_name' in y_2.name), True)
+            self.assertEqual(('nuclear_name' in y_3.name), True)
 
     def test_errors(self):
         with base.program_guard(base.Program(), base.Program()):