agedi.diffusion.distributions.categorical
=========================================

.. py:module:: agedi.diffusion.distributions.categorical


Classes
-------

.. autoapisummary::

   agedi.diffusion.distributions.categorical.Categorical


Module Contents
---------------

.. py:class:: Categorical

   Bases: :py:obj:`agedi.diffusion.distributions.Distribution`


   Categorical Distribution

   Implements hard sampling using the Gumbel-Max trick.



   .. py:method:: _sample(probs: torch.Tensor) -> torch.Tensor

      Sample from the categorical distribution where
      probabilites define the likelihood of mu value
      to be set to the masked, 0, value

      :param probs: The probabilities of each category
      :type probs: torch.Tensor

      :returns: Sampled tensor
      :rtype: torch.Tensor