code: def _rescale_to_unit_diagonals(mat):
      
      d = np.sqrt(np.diag(mat))
      mat /= d
      mat /= d[:, np.newaxis]
      return mat
docs: Rescale matrix to have unit diagonals.
  
      Note: Call only after diagonal dominance is ensured.