def _params(self, inputs):
    """Extracts the mean and std parameters from the inputs."""
    if inputs.shape[-1] != self.n_inputs:
      raise ValueError(
          'Invalid distribution parametrization - expected {} parameters, '
          'got {}. Input shape: {}.'.format(
              self.n_inputs, inputs.shape[-1], inputs.shape
          )
      )
    n_dims = self._n_dims
    # Split the distribution inputs into two parts: mean and std.
    mean = inputs[..., :n_dims]
    if self._learn_std is not None:
      std = inputs[..., n_dims:]
      # Std is non-negative, so let's softplus it.
      std = tl.Softplus()(std + self._std)
    else:
      std = self._std
    # In case of constant or shared std, upsample it to the same dimensionality
    # as the means.
    std = jnp.broadcast_to(std, mean.shape)
    return (mean, std)
Extracts the mean and std parameters from the inputs.
trax/rl/distributions.py