def __init__(self, n_categories, shape=()):
    """Initializes Categorical distribution.

    Args:
      n_categories (int): Number of categories.
      shape (tuple): Shape of the sample.
    """
    self._n_categories = n_categories
    self._shape = shape
Initializes Categorical distribution.

    Args:
      n_categories (int): Number of categories.
      shape (tuple): Shape of the sample.
    
trax/rl/distributions.py