Fix paddle.distribution.LKJCholesky.sample() when concentration is 1-D 易用性提升 (#68922)

NKNaN · web-flow · commit 01a04d7af5e7 · 2024-11-12T11:58:39.000+08:00
* fix

* update for 0-D

* fix for ci-kunlun-r200
diff --git a/python/paddle/distribution/lkj_cholesky.py b/python/paddle/distribution/lkj_cholesky.py
@@ -75,7 +75,7 @@ def vec_to_tril_matrix(
     shape0 = flatten_shape // last_dim
     output_shape = (
         *sample_shape,
-        shape0 // reduce(operator.mul, sample_shape),
+        shape0 // reduce(operator.mul, sample_shape, 1),
         dim,
         dim,
     )
@@ -120,6 +120,7 @@ def tril_matrix_to_vec(mat: Tensor, diag: int = 0) -> Tensor:
     out_shape += (vec_len,)
 
     # Use the mask to index the lower triangular elements from the input matrix
+    tril_mask = paddle.broadcast_to(tril_mask, mat.shape)
     vec = paddle.masked_select(mat, tril_mask).reshape(out_shape)
     return vec
 
@@ -192,13 +193,6 @@ def __init__(
             if not paddle.all(self.concentration > 0):
                 raise ValueError("The arg of `concentration` must be positive.")
 
-        self.concentration = concentration
-        if isinstance(self.concentration, float):
-            self.concentration = (self.concentration,)
-
-        if not isinstance(self.concentration, paddle.Tensor):
-            self.concentration = paddle.to_tensor(self.concentration)
-
         self.sample_method = sample_method
         batch_shape = self.concentration.shape
         event_shape = (dim, dim)
@@ -284,8 +278,8 @@ def _cvine(self, sample_shape: Sequence[int]) -> Tensor:
 
         # Construct the lower triangular matrix from the partial correlations
         last_dim = self.dim * (self.dim - 1) // 2
-        flatten_shape = last_dim * reduce(operator.mul, sample_shape)
-        if self.concentration.shape != ():
+        flatten_shape = last_dim * reduce(operator.mul, sample_shape, 1)
+        if len(self.concentration.shape) != 0:
             flatten_shape *= self.concentration.shape[-1]
 
         partial_correlation = partial_correlation.reshape((flatten_shape,))
@@ -320,32 +314,20 @@ def _cvine(self, sample_shape: Sequence[int]) -> Tensor:
             partial_correlation.shape[-2], partial_correlation.shape[-1]
         )
         r = r * z1m_cumprod_sqrt_shifted
-        if sample_shape == (1,):
-            r = r.reshape((flatten_shape // last_dim, self.dim, self.dim))
         return r
 
     def sample(self, sample_shape: Sequence[int] = ()) -> Tensor:
         """Generate a sample using the specified sampling method."""
         if not isinstance(sample_shape, Sequence):
             raise TypeError('sample shape must be Sequence object.')
 
-        # for paddle.static, U need to set sample_shape
-        if sample_shape == ():
-            sample_shape = (1,)
         if self.sample_method == "onion":
             res = self._onion(sample_shape)
         else:
             res = self._cvine(sample_shape)
 
-        output_shape = []
-        if sample_shape != (1,):
-            output_shape = list(sample_shape)
-
-        if tuple(self.concentration.shape) != () and tuple(
-            self.concentration.shape
-        ) != (1,):
-            output_shape.extend(self.concentration.shape)
-
+        output_shape = list(sample_shape)
+        output_shape.extend(self.concentration.shape)
         output_shape.extend([self.dim, self.dim])
 
         return res.reshape(output_shape)
diff --git a/test/distribution/test_distribution_lkj_cholesky.py b/test/distribution/test_distribution_lkj_cholesky.py
@@ -33,6 +33,15 @@
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'concentration'),
     [
+        (
+            'zero-dim',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ).reshape([]),
+        ),
         (
             'one-dim',
             parameterize.xrand(
@@ -42,14 +51,23 @@
                 min=0,
             ),
         ),
+        (
+            'one-dim2',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ),
+        ),
     ],
 )
 class TestLKJCholeskyShape(unittest.TestCase):
     def gen_cases(self):
-        extra_shape = (
-            len(self.concentration),
-            self._paddle_lkj_cholesky.dim,
-            self._paddle_lkj_cholesky.dim,
+        extra_shape = []
+        extra_shape.extend(self.concentration.shape)
+        extra_shape.extend(
+            [self._paddle_lkj_cholesky.dim, self._paddle_lkj_cholesky.dim]
         )
         cases = [
             {
diff --git a/test/distribution/test_distribution_lkj_cholesky_static.py b/test/distribution/test_distribution_lkj_cholesky_static.py
@@ -35,6 +35,15 @@
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'concentration'),
     [
+        (
+            'zerp-dim',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ).reshape([]),
+        ),
         (
             'one-dim',
             parameterize.xrand(
@@ -44,6 +53,15 @@
                 min=0,
             ),
         ),
+        (
+            'one-dim2',
+            parameterize.xrand(
+                (1,),
+                dtype='float64',
+                max=1.0,
+                min=0,
+            ),
+        ),
     ],
 )
 class TestLKJCholeskyShape(unittest.TestCase):
@@ -70,10 +88,10 @@ def setUp(self):
             }
 
     def gen_cases(self):
-        extra_shape = (
-            len(self.concentration),
-            self._paddle_lkj_cholesky.dim,
-            self._paddle_lkj_cholesky.dim,
+        extra_shape = []
+        extra_shape.extend(self.concentration.shape)
+        extra_shape.extend(
+            [self._paddle_lkj_cholesky.dim, self._paddle_lkj_cholesky.dim]
         )
         cases = [
             {
