def exponent_for_preconditioner(self):
        """Returns exponent to use for inverse-pth root M^{-1/p}."""
        should_preconditioned_dims = self.should_precondition_dims()
        num_preconditioners = sum(should_preconditioned_dims)
        return 2 * num_preconditioners
Returns exponent to use for inverse-pth root M^{-1/p}.
tools/train/scalable_shampoo/distributed_shampoo.py