Fix distribution API with 0-D input parameters 易用性提升 -part1 (#69090)

NKNaN · web-flow · commit 563317aa580b · 2024-11-07T12:21:15.000+08:00
* update

* update test

* fix test
diff --git a/python/paddle/distribution/bernoulli.py b/python/paddle/distribution/bernoulli.py
@@ -146,11 +146,11 @@ def variance(self) -> Tensor:
         """
         return paddle.multiply(self.probs, (1 - self.probs))
 
-    def sample(self, shape: Sequence[int]) -> Tensor:
+    def sample(self, shape: Sequence[int] = ()) -> Tensor:
         """Sample from Bernoulli distribution.
 
         Args:
-            shape (Sequence[int]): Sample shape.
+            shape (Sequence[int], optional): Sample shape.
 
         Returns:
             Tensor: Sampled data with shape `sample_shape` + `batch_shape` + `event_shape`.
@@ -193,7 +193,9 @@ def sample(self, shape: Sequence[int]) -> Tensor:
         with paddle.no_grad():
             return paddle.bernoulli(self.probs.expand(shape), name=name)
 
-    def rsample(self, shape: Sequence[int], temperature: float = 1.0) -> Tensor:
+    def rsample(
+        self, shape: Sequence[int] = (), temperature: float = 1.0
+    ) -> Tensor:
         """Sample from Bernoulli distribution (reparameterized).
 
         The `rsample` is a continuously approximate of Bernoulli distribution reparameterized sample method.
@@ -204,7 +206,7 @@ def rsample(self, shape: Sequence[int], temperature: float = 1.0) -> Tensor:
             `rsample` need to be followed by a `sigmoid`, which converts samples' value to unit interval (0, 1).
 
         Args:
-            shape (Sequence[int]): Sample shape.
+            shape (Sequence[int], optional): Sample shape.
             temperature (float): temperature for rsample, must be positive.
 
         Returns:
diff --git a/python/paddle/distribution/binomial.py b/python/paddle/distribution/binomial.py
@@ -84,10 +84,7 @@ def __init__(
         self.dtype = paddle.get_default_dtype()
         self.total_count, self.probs = self._to_tensor(total_count, probs)
 
-        if self.total_count.shape == []:
-            batch_shape = (1,)
-        else:
-            batch_shape = self.total_count.shape
+        batch_shape = self.total_count.shape
         super().__init__(batch_shape)
 
     def _to_tensor(
diff --git a/python/paddle/distribution/categorical.py b/python/paddle/distribution/categorical.py
@@ -148,11 +148,11 @@ def __init__(
         dist_sum = paddle.sum(self.logits, axis=-1, keepdim=True)
         self._prob = self.logits / dist_sum
 
-    def sample(self, shape: Sequence[int]) -> Tensor:
+    def sample(self, shape: Sequence[int] = ()) -> Tensor:
         """Generate samples of the specified shape.
 
         Args:
-            shape (list): Shape of the generated samples.
+            shape (list, optional): Shape of the generated samples.
 
         Returns:
             Tensor: A tensor with prepended dimensions shape.
diff --git a/python/paddle/distribution/cauchy.py b/python/paddle/distribution/cauchy.py
@@ -121,14 +121,16 @@ def stddev(self) -> Never:
         """Standard Deviation of Cauchy distribution."""
         raise ValueError("Cauchy distribution has no stddev.")
 
-    def sample(self, shape: Sequence[int], name: str | None = None) -> Tensor:
+    def sample(
+        self, shape: Sequence[int] = (), name: str | None = None
+    ) -> Tensor:
         """Sample from Cauchy distribution.
 
         Note:
             `sample` method has no grad, if you want so, please use `rsample` instead.
 
         Args:
-            shape (Sequence[int]): Sample shape.
+            shape (Sequence[int], optional): Sample shape.
             name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
         Returns:
@@ -169,11 +171,13 @@ def sample(self, shape: Sequence[int], name: str | None = None) -> Tensor:
         with paddle.no_grad():
             return self.rsample(shape, name)
 
-    def rsample(self, shape: Sequence[int], name: str | None = None) -> Tensor:
+    def rsample(
+        self, shape: Sequence[int] = (), name: str | None = None
+    ) -> Tensor:
         """Sample from Cauchy distribution (reparameterized).
 
         Args:
-            shape (Sequence[int]): Sample shape.
+            shape (Sequence[int], optional): Sample shape.
             name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
         Returns:
diff --git a/python/paddle/distribution/continuous_bernoulli.py b/python/paddle/distribution/continuous_bernoulli.py
@@ -118,10 +118,7 @@ def __init__(
         eps_prob = paddle.finfo(self.probs.dtype).eps
         self.probs = paddle.clip(self.probs, min=eps_prob, max=1 - eps_prob)
 
-        if self.probs.shape == []:
-            batch_shape = (1,)
-        else:
-            batch_shape = self.probs.shape
+        batch_shape = self.probs.shape
         super().__init__(batch_shape)
 
     def _to_tensor(self, probs: float | Tensor) -> Tensor:
diff --git a/test/distribution/test_distribution_binomial.py b/test/distribution/test_distribution_binomial.py
@@ -27,6 +27,11 @@
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
     [
+        (
+            'zero-dim',
+            np.array(1000),
+            np.array(0.6),
+        ),
         (
             'one-dim',
             1000,
@@ -127,7 +132,7 @@ def _np_entropy(self):
 class TestBinomialProbs(unittest.TestCase):
     def setUp(self):
         self._dist = Binomial(
-            total_count=self.total_count,
+            total_count=paddle.to_tensor(self.total_count),
             probs=paddle.to_tensor(self.probs),
         )
 
diff --git a/test/distribution/test_distribution_binomial_static.py b/test/distribution/test_distribution_binomial_static.py
@@ -32,6 +32,11 @@
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
     [
+        (
+            'zero-dim',
+            np.array(1000),
+            np.array(0.6),
+        ),
         (
             'one-dim',
             np.array([1000]),
@@ -127,6 +132,12 @@ def _np_entropy(self):
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),
     [
+        (
+            'zero-dim',
+            np.array(10),
+            np.array(0.6).astype('float64'),
+            np.array([2.0, 3.0, 5.0]).astype('float64'),
+        ),
         (
             'value-same-shape',
             np.array([10]).astype('int64'),
diff --git a/test/distribution/test_distribution_continuous_bernoulli.py b/test/distribution/test_distribution_continuous_bernoulli.py
@@ -241,6 +241,11 @@ def test_sample(self):
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'probs', 'value'),
     [
+        (
+            'zero-dim',
+            np.array(0.3).astype("float32"),
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),
+        ),
         (
             'value-same-shape',
             parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),
@@ -297,6 +302,11 @@ def test_icdf(self):
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'p_1', 'p_2'),
     [
+        (
+            'zero-dim',
+            np.array(0.2).astype("float32"),
+            np.array(0.4).astype("float32"),
+        ),
         (
             'one-dim',
             parameterize.xrand((1,), min=0.0, max=1.0).astype("float32"),
diff --git a/test/distribution/test_distribution_continuous_bernoulli_static.py b/test/distribution/test_distribution_continuous_bernoulli_static.py
@@ -159,6 +159,10 @@ def np_kl_divergence(self, other):
 @parameterize.parameterize_cls(
     (parameterize.TEST_CASE_NAME, 'probs'),
     [
+        (
+            'zero-dim',
+            np.array(0.7).astype("float32"),
+        ),
         (
             'multi-dim',
             parameterize.xrand((1, 3), min=0.0, max=1.0).astype("float32"),

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,11 @@`
`27`	`27`	`@parameterize.parameterize_cls(`
`28`	`28`	`(parameterize.TEST_CASE_NAME, 'total_count', 'probs'),`
`29`	`29`	`[`
	`30`	`+ (`
	`31`	`+ 'zero-dim',`
	`32`	`+ np.array(1000),`
	`33`	`+ np.array(0.6),`
	`34`	`+ ),`
`30`	`35`	`(`
`31`	`36`	`'one-dim',`
`32`	`37`	`1000,`
`@@ -127,7 +132,7 @@ def _np_entropy(self):`
`127`	`132`	`class TestBinomialProbs(unittest.TestCase):`
`128`	`133`	`def setUp(self):`
`129`	`134`	`self._dist = Binomial(`
`130`		`- total_count=self.total_count,`
	`135`	`+ total_count=paddle.to_tensor(self.total_count),`
`131`	`136`	`probs=paddle.to_tensor(self.probs),`
`132`	`137`	`)`
`133`	`138`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,11 @@`
`32`	`32`	`@parameterize.parameterize_cls(`
`33`	`33`	`(parameterize.TEST_CASE_NAME, 'total_count', 'probs'),`
`34`	`34`	`[`
	`35`	`+ (`
	`36`	`+ 'zero-dim',`
	`37`	`+ np.array(1000),`
	`38`	`+ np.array(0.6),`
	`39`	`+ ),`
`35`	`40`	`(`
`36`	`41`	`'one-dim',`
`37`	`42`	`np.array([1000]),`
`@@ -127,6 +132,12 @@ def _np_entropy(self):`
`127`	`132`	`@parameterize.parameterize_cls(`
`128`	`133`	`(parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),`
`129`	`134`	`[`
	`135`	`+ (`
	`136`	`+ 'zero-dim',`
	`137`	`+ np.array(10),`
	`138`	`+ np.array(0.6).astype('float64'),`
	`139`	`+ np.array([2.0, 3.0, 5.0]).astype('float64'),`
	`140`	`+ ),`
`130`	`141`	`(`
`131`	`142`	`'value-same-shape',`
`132`	`143`	`np.array([10]).astype('int64'),`
Original file line number	Diff line number	Diff line change
`@@ -241,6 +241,11 @@ def test_sample(self):`
`241`	`241`	`@parameterize.parameterize_cls(`
`242`	`242`	`(parameterize.TEST_CASE_NAME, 'probs', 'value'),`
`243`	`243`	`[`
	`244`	`+ (`
	`245`	`+ 'zero-dim',`
	`246`	`+ np.array(0.3).astype("float32"),`
	`247`	`+ parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),`
	`248`	`+ ),`
`244`	`249`	`(`
`245`	`250`	`'value-same-shape',`
`246`	`251`	`parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),`
`@@ -297,6 +302,11 @@ def test_icdf(self):`
`297`	`302`	`@parameterize.parameterize_cls(`
`298`	`303`	`(parameterize.TEST_CASE_NAME, 'p_1', 'p_2'),`
`299`	`304`	`[`
	`305`	`+ (`
	`306`	`+ 'zero-dim',`
	`307`	`+ np.array(0.2).astype("float32"),`
	`308`	`+ np.array(0.4).astype("float32"),`
	`309`	`+ ),`
`300`	`310`	`(`
`301`	`311`	`'one-dim',`
`302`	`312`	`parameterize.xrand((1,), min=0.0, max=1.0).astype("float32"),`
Original file line number	Diff line number	Diff line change
`@@ -159,6 +159,10 @@ def np_kl_divergence(self, other):`
`159`	`159`	`@parameterize.parameterize_cls(`
`160`	`160`	`(parameterize.TEST_CASE_NAME, 'probs'),`
`161`	`161`	`[`
	`162`	`+ (`
	`163`	`+ 'zero-dim',`
	`164`	`+ np.array(0.7).astype("float32"),`
	`165`	`+ ),`
`162`	`166`	`(`
`163`	`167`	`'multi-dim',`
`164`	`168`	`parameterize.xrand((1, 3), min=0.0, max=1.0).astype("float32"),`