wehub-resource-sync
154 lines
5.2 KiB
Python
154 lines
5.2 KiB
Python
# 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
|