docs: def loss_l2(self, l2=0):
          """L2 loss centered around mu_init, scaled optionally per-source.
  
          In other words, diagonal Tikhonov regularization,
              ||D(\mu-\mu_{init})||_2^2
          where D is diagonal.
  
          Args:
              - l2: A float or np.array representing the per-source regularization
                  strengths to use
          """
          if isinstance(l2, (int, float)):
              D = l2 * torch.eye(self.d)
          else:
              D = torch.diag(torch.from_numpy(l2))
  
          # Note that mu is a matrix and this is the *Frobenius norm*
          return torch.norm(D @ (self.mu - self.mu_init)) ** 2