chore: import upstream snapshot with attribution

python/paddle/compat/distributions/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2026 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.
+
+from __future__ import annotations
+
+from . import categorical as categorical
+from .categorical import Categorical
+
+__all__ = [
+    'Categorical',
+]

python/paddle/compat/distributions/categorical.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2026 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.
+
+from __future__ import annotations
+
+import numpy as np
+
+import paddle
+from paddle.distribution import constraint, distribution
+from paddle.tensor import multinomial
+
+
+class Categorical(distribution.Distribution):
+    arg_constraints = {
+        "probs": constraint.simplex,
+        "logits": constraint.real_vector,
+    }
+    has_enumerate_support = True
+
+    def __init__(
+        self,
+        probs=None,
+        logits=None,
+        validate_args: bool | None = None,
+    ) -> None:
+        if (probs is None) == (logits is None):
+            raise ValueError(
+                "Either `probs` or `logits` must be specified, but not both."
+            )
+        if probs is not None:
+            if probs.dim() < 1:
+                raise ValueError(
+                    "`probs` parameter must be at least one-dimensional."
+                )
+            self._probs = probs / probs.sum(-1, keepdim=True)
+            self._logits = None
+            self._param = self._probs
+        else:
+            if logits.dim() < 1:
+                raise ValueError(
+                    "`logits` parameter must be at least one-dimensional."
+                )
+            self._logits = logits - paddle.logsumexp(
+                logits, axis=-1, keepdim=True
+            )
+            self._probs = None
+            self._param = self._logits
+
+        self._num_events = self._param.shape[-1]
+        batch_shape = (
+            tuple(self._param.shape[:-1]) if self._param.dim() > 1 else ()
+        )
+        super().__init__(batch_shape, validate_args=validate_args)
+
+    def expand(self, batch_shape, _instance=None):
+        new = (
+            self.__class__.__new__(self.__class__)
+            if _instance is None
+            else _instance
+        )
+        batch_shape = tuple(batch_shape)
+        param_shape = (*batch_shape, self._num_events)
+        new._probs = (
+            self._probs.expand(param_shape) if self._probs is not None else None
+        )
+        new._logits = (
+            self._logits.expand(param_shape)
+            if self._logits is not None
+            else None
+        )
+        new._param = new._logits if new._logits is not None else new._probs
+        new._num_events = self._num_events
+        super(Categorical, new).__init__(batch_shape, validate_args=False)
+        new._validate_args_enabled = self._validate_args_enabled
+        return new
+
+    @property
+    def support(self):
+        return constraint.integer_interval(0, self._num_events - 1)
+
+    @property
+    def logits(self):
+        if self._logits is None:
+            eps = paddle.finfo(self.probs.dtype).eps
+            probs = paddle.clip(self.probs, min=eps, max=1 - eps)
+            self._logits = paddle.log(probs)
+        return self._logits
+
+    @property
+    def probs(self):
+        if self._probs is None:
+            self._probs = paddle.nn.functional.softmax(self.logits, axis=-1)
+        return self._probs
+
+    @property
+    def param_shape(self):
+        return self._param.shape
+
+    @property
+    def mean(self):
+        return paddle.full_like(self.probs[..., 0], float('nan'))
+
+    @property
+    def mode(self):
+        return paddle.argmax(self.probs, axis=-1)
+
+    @property
+    def variance(self):
+        return paddle.full_like(self.probs[..., 0], float('nan'))
+
+    def sample(self, sample_shape=()):
+        sample_shape = tuple(sample_shape)
+        probs_2d = self.probs.reshape([-1, self._num_events])
+        samples_2d = multinomial(probs_2d, int(np.prod(sample_shape)), True).T
+        return samples_2d.reshape(self._extend_shape(sample_shape))
+
+    def log_prob(self, value):
+        if self._validate_args_enabled and paddle.in_dynamic_mode():
+            self._validate_sample(value)
+        value = paddle.cast(value, dtype='int64').unsqueeze(-1)
+        log_pmf = self.logits
+        output_shape = paddle.broadcast_shape(value.shape, log_pmf.shape)
+        value = paddle.broadcast_to(value, [*output_shape[:-1], 1])
+        log_pmf = paddle.broadcast_to(log_pmf, output_shape)
+        return paddle.take_along_axis(log_pmf, value, axis=-1).squeeze(-1)
+
+    def entropy(self):
+        min_real = paddle.finfo(self.logits.dtype).min
+        logits = paddle.clip(self.logits, min=min_real)
+        p_log_p = logits * self.probs
+        return -p_log_p.sum(-1)
+
+    def enumerate_support(self, expand=True):
+        values = paddle.arange(self._num_events, dtype='int64')
+        values = values.reshape(
+            [self._num_events] + [1] * len(self._batch_shape)
+        )
+        if expand:
+            values = paddle.broadcast_to(
+                values, [self._num_events, *self._batch_shape]
+            )
+        return values