    def __init__(
        self,
        num_groups: int,
        whitening_limit: float,
        prob: Union[float, Tuple[float, float]],
        grad_scale: float,
    ):
        """
        Args:
          num_groups: the number of groups to divide the channel dim into before
            whitening.  We will attempt to make the feature covariance
            within each group, after mean subtraction, as "white" as possible,
            while having the same trace across all groups.
         whitening_limit: a value greater than 1.0, that dictates how much
           freedom we have to violate the constraints.  1.0 would mean perfectly
           white, with exactly the same trace across groups; larger values
           give more freedom.  E.g. 2.0.
         prob: the probability with which we apply the gradient modification
           (also affects the grad scale).  May be supplied as a float,
           or as a pair (min_prob, max_prob)

          grad_scale: determines the scale on the gradient term from this object,
            relative to the rest of the gradient on the attention weights.
            E.g. 0.02 (you may want to use smaller values than this if prob is large)
        """
        super(Whiten, self).__init__()
        assert num_groups >= 1
        assert whitening_limit >= 1
        assert grad_scale >= 0
        self.num_groups = num_groups
        self.whitening_limit = whitening_limit
        if isinstance(prob, float):
            assert 0 < prob <= 1
            self.prob = prob
        else:
            (self.min_prob, self.max_prob) = prob
            assert 0 < self.min_prob < self.max_prob <= 1
            self.prob = self.max_prob

        self.grad_scale = grad_scale
Args:
num_groups: the number of groups to divide the channel dim into before
whitening.  We will attempt to make the feature covariance
within each group, after mean subtraction, as "white" as possible,
while having the same trace across all groups.
whitening_limit: a value greater than 1.0, that dictates how much
freedom we have to violate the constraints.  1.0 would mean perfectly
white, with exactly the same trace across groups; larger values
give more freedom.  E.g. 2.0.
prob: the probability with which we apply the gradient modification
(also affects the grad scale).  May be supplied as a float,
or as a pair (min_prob, max_prob)

grad_scale: determines the scale on the gradient term from this object,
relative to the rest of the gradient on the attention weights.
E.g. 0.02 (you may want to use smaller values than this if prob is large)
modules/general/scaling.py