    def forward(self, x: Tensor) -> Tensor:
        """
        In the forward pass, this function just returns the input unmodified.
        In the backward pass, it will modify the gradients to ensure that the
        distribution in each group has close to (lambda times I) as the covariance
        after mean subtraction, with the same lambda across groups.
        For whitening_limit > 1, there will be more freedom to violate this
        constraint.

        Args:
           x: the input of shape (*, num_channels)

        Returns:
            x, unmodified.   You should make sure
        you use the returned value, or the graph will be freed
        and nothing will happen in backprop.
        """
        if not x.requires_grad or random.random() > self.prob or self.grad_scale == 0:
            return _no_op(x)
        else:
            if hasattr(self, "min_prob") and random.random() < 0.25:
                # occasionally switch between min_prob and max_prob, based on whether
                # we are above or below the threshold.
                if (
                    _whitening_metric(x.to(torch.float32), self.num_groups)
                    > self.whitening_limit
                ):
                    # there would be a change to the grad.
                    self.prob = self.max_prob
                else:
                    self.prob = self.min_prob

            return WhiteningPenaltyFunction.apply(
                x, self.num_groups, self.whitening_limit, self.grad_scale
            )
In the forward pass, this function just returns the input unmodified.
In the backward pass, it will modify the gradients to ensure that the
distribution in each group has close to (lambda times I) as the covariance
after mean subtraction, with the same lambda across groups.
For whitening_limit > 1, there will be more freedom to violate this
constraint.

Args:
x: the input of shape (*, num_channels)

Returns:
x, unmodified.   You should make sure
you use the returned value, or the graph will be freed
and nothing will happen in backprop.
modules/general/scaling.py