def __init__(self, shape=(), std=1.0, learn_std=None):
    """Initializes Gaussian distribution.

    Args:
      shape (tuple): Shape of the sample.
      std (float): Standard deviation, shared across the whole sample.
      learn_std (str or None): How to learn the standard deviation - 'shared'
        to have a single, shared std parameter, or 'separate' to have separate
        parameters for each dimension.
    """
    self._shape = shape
    self._std = std
    self._learn_std = learn_std
Initializes Gaussian distribution.

    Args:
      shape (tuple): Shape of the sample.
      std (float): Standard deviation, shared across the whole sample.
      learn_std (str or None): How to learn the standard deviation - 'shared'
        to have a single, shared std parameter, or 'separate' to have separate
        parameters for each dimension.
    
trax/rl/distributions.py