# Copyright (c) 2021 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 typing import TYPE_CHECKING, Union

import numpy as np

import paddle

from paddle.base.data_feeder import check_type, convert_dtype

from paddle.base.framework import Variable

from paddle.distribution import distribution

from paddle.framework import in_dynamic_mode

from paddle.tensor import multinomial

if TYPE_CHECKING:

from collections.abc import Sequence

import numpy.typing as npt

from typing_extensions import TypeAlias

from paddle import Tensor

from paddle._typing import NestedSequence

from paddle._typing.dtype_like import _DTypeLiteral

_CategoricalBoundary: TypeAlias = Union[

Sequence[float],

NestedSequence[float],

npt.NDArray[Union[np.float32, np.float64]],

Tensor,

]

class Categorical(distribution.Distribution):

r"""

Categorical distribution is a discrete probability distribution that

describes the possible results of a random variable that can take on

one of K possible categories, with the probability of each category

separately specified.

The probability mass function (pmf) is:

.. math::

pmf(k; p_i) = \prod_{i=1}^{k} p_i^{[x=i]}

In the above equation:

* :math:`[x=i]` : it evaluates to 1 if :math:`x==i` , 0 otherwise.

Args:

logits(list|tuple|numpy.ndarray|Tensor): The logits input of categorical distribution. The data type is float32 or float64.

name(str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> paddle.seed(200) # on CPU device

>>> y = paddle.rand([6])

>>> print(y)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.77663314, 0.90824795, 0.15685187, 0.04279523, 0.34468332, 0.79557180])

>>> cat = Categorical(x)

>>> cat2 = Categorical(y)

>>> paddle.seed(1000) # on CPU device

>>> print(cat.sample([2, 3]))

Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,

[[0, 1, 5],

[3, 4, 5]])

>>> print(cat.entropy())

Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,

1.77528250)

>>> print(cat.kl_divergence(cat2))

Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.07195196])

>>> value = paddle.to_tensor([2, 1, 3])

>>> print(cat.probs(value))

Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.00608027, 0.10829761, 0.26965630])

>>> print(cat.log_prob(value))

Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,

[-5.10270691, -2.22287226, -1.31060708])

"""

logits: Tensor

dtype: _DTypeLiteral

def __init__(

self,

logits: _CategoricalBoundary,

name: str | None = None,

) -> None:

"""

Args:

logits(list|tuple|numpy.ndarray|Tensor): The logits input of categorical distribution. The data type is float32 or float64.

name(str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.

"""

if not in_dynamic_mode():

check_type(

logits,

'logits',

(np.ndarray, Variable, paddle.pir.Value, list, tuple),

'Categorical',

)

self.name = name if name is not None else 'Categorical'

self.dtype = 'float32'

if self._validate_args(logits):

self.logits = logits

self.dtype = convert_dtype(logits.dtype)

else:

if isinstance(logits, np.ndarray) and str(logits.dtype) in [

'float32',

'float64',

]:

self.dtype = convert_dtype(logits.dtype)

self.logits = self._to_tensor(logits)[0]

if self.dtype != convert_dtype(self.logits.dtype):

self.logits = paddle.cast(self.logits, dtype=self.dtype)

dist_sum = paddle.sum(self.logits, axis=-1, keepdim=True)

self._prob = self.logits / dist_sum

def sample(self, shape: Sequence[int] = []) -> Tensor:

"""Generate samples of the specified shape.

Args:

shape (Sequence[int], optional): Shape of the generated samples.

Returns:

Tensor: A tensor with prepended dimensions shape.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> # doctest: +SKIP('Random output')

>>> cat = Categorical(x)

>>> paddle.seed(1000) # on CPU device

>>> print(cat.sample([2, 3]))

Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,

[[0, 1, 5],

[3, 4, 5]])

"""

name = self.name + '_sample'

if not in_dynamic_mode():

check_type(shape, 'shape', (list, tuple), 'sample')

num_samples = np.prod(np.array(shape))

logits_shape = list(self.logits.shape)

if len(logits_shape) > 1:

sample_shape = shape + logits_shape[:-1]

logits = paddle.reshape(

self.logits, [np.prod(logits_shape[:-1]), logits_shape[-1]]

)

else:

sample_shape = shape

logits = self.logits

sample_index = multinomial(

self._logits_to_probs(logits), num_samples, True

)

# multinomial sample shape is (logits.shape[:-1], num_samples), need to

# transpose to (num_samples, logits.shape[:-1])

permute = list(range(sample_index.dim()))

permute.insert(0, permute.pop(-1))

sample_index = sample_index.transpose(permute)

return paddle.reshape(sample_index, sample_shape, name=name)

def kl_divergence(self, other: Categorical) -> Tensor:

"""The KL-divergence between two Categorical distributions.

Args:

other (Categorical): instance of Categorical. The data type is float32.

Returns:

Tensor: kl-divergence between two Categorical distributions.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> paddle.seed(200) # on CPU device

>>> y = paddle.rand([6])

>>> print(y)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.77663314, 0.90824795, 0.15685187, 0.04279523, 0.34468332, 0.79557180])

>>> cat = Categorical(x)

>>> cat2 = Categorical(y)

>>> print(cat.kl_divergence(cat2))

Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.07195196])

"""

name = self.name + '_kl_divergence'

if not in_dynamic_mode():

check_type(other, 'other', Categorical, 'kl_divergence')

logits = self.logits - paddle.max(self.logits, axis=-1, keepdim=True)

other_logits = other.logits - paddle.max(

other.logits, axis=-1, keepdim=True

)

e_logits = paddle.exp(logits)

other_e_logits = paddle.exp(other_logits)

z = paddle.sum(e_logits, axis=-1, keepdim=True)

other_z = paddle.sum(other_e_logits, axis=-1, keepdim=True)

prob = e_logits / z

kl = paddle.sum(

prob

* (logits - paddle.log(z) - other_logits + paddle.log(other_z)),

axis=-1,

keepdim=True,

name=name,

)

return kl

def entropy(self) -> Tensor:

"""Shannon entropy in nats.

Returns:

Tensor: Shannon entropy of Categorical distribution. The data type is float32.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> cat = Categorical(x)

>>> print(cat.entropy())

Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,

1.77528250)

"""

name = self.name + '_entropy'

logits = self.logits - paddle.max(self.logits, axis=-1, keepdim=True)

e_logits = paddle.exp(logits)

z = paddle.sum(e_logits, axis=-1, keepdim=True)

prob = e_logits / z

neg_entropy = paddle.sum(prob * (logits - paddle.log(z)), axis=-1)

entropy = paddle.scale(neg_entropy, scale=-1.0, name=name)

return entropy

def probs(self, value: Tensor) -> Tensor:

"""Probabilities of the given category (``value``).

If ``logits`` is 2-D or higher dimension, the last dimension will be regarded as

category, and the others represents the different distributions.

At the same time, if ``value`` is 1-D Tensor, ``value`` will be broadcast to the

same number of distributions as ``logits``.

If ``value`` is not 1-D Tensor, ``value`` should have the same number distributions

with ``logits. That is, ``value[:-1] = logits[:-1]``.

Args:

value (Tensor): The input tensor represents the selected category index.

Returns:

Tensor: probability according to the category index.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> cat = Categorical(x)

>>> value = paddle.to_tensor([2, 1, 3])

>>> print(cat.probs(value))

Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.00608027, 0.10829761, 0.26965630])

"""

name = self.name + '_probs'

if len(self._prob.shape) == 1: # batch_shape is empty

return paddle.gather(

self._prob, value.reshape([-1], name=name), name=name

).reshape(value.shape, name=name)

else:

if len(value.shape) == 1:

return paddle.take_along_axis(

self._prob,

paddle.reshape(

value,

(len(self._prob.shape) - 1) * [1] + [-1],

name=name,

),

axis=-1,

)

else:

return paddle.take_along_axis(self._prob, value, axis=-1)

def log_prob(self, value: Tensor) -> Tensor:

"""Log probabilities of the given category. Refer to ``probs`` method.

Args:

value (Tensor): The input tensor represents the selected category index.

Returns:

Tensor: Log probability.

Examples:

.. code-block:: pycon

>>> import paddle

>>> from paddle.distribution import Categorical

>>> paddle.seed(100) # on CPU device

>>> x = paddle.rand([6])

>>> print(x)

Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,

[0.55355281, 0.20714243, 0.01162981, 0.51577556, 0.36369765, 0.26091650])

>>> cat = Categorical(x)

>>> value = paddle.to_tensor([2, 1, 3])

>>> print(cat.log_prob(value))

Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,

[-5.10270691, -2.22287226, -1.31060708])

"""

name = self.name + '_log_prob'

return paddle.log(self.probs(value), name=name)