def METHOD_NAME(mat_a, mat_a_size, iter_count=100, ridge_epsilon=1e-4):
  """Iterative method to get matrix square root.
  Stable iterations for the matrix square root, Nicholas J. Higham
  Page 231, Eq 2.6b
  http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.6.8799&rep=rep1&type=pdf
  Args:
    mat_a: the symmetric PSD matrix whose matrix square root be computed
    mat_a_size: size of mat_a.
    iter_count: Maximum number of iterations.
    ridge_epsilon: Ridge epsilon added to make the matrix positive definite.
  Returns:
    mat_a^0.5
  """
  def _iter_condition(i, unused_mat_y, unused_old_mat_y, unused_mat_z,
                      unused_old_mat_z, err, old_err):
    # This method require that we check for divergence every step.
    return math_ops.logical_and(i < iter_count, err < old_err)
  def _iter_body(i, mat_y, unused_old_mat_y, mat_z, unused_old_mat_z, err,
                 unused_old_err):
    current_iterate = 0.5 * (3.0 * identity - math_ops.matmul(mat_z, mat_y))
    current_mat_y = math_ops.matmul(mat_y, current_iterate)
    current_mat_z = math_ops.matmul(current_iterate, mat_z)
    # Compute the error in approximation.
    mat_sqrt_a = current_mat_y * math_ops.sqrt(norm)
    mat_a_approx = math_ops.matmul(mat_sqrt_a, mat_sqrt_a)
    residual = mat_a - mat_a_approx
    current_err = math_ops.sqrt(math_ops.reduce_sum(residual * residual)) / norm
    return i + 1, current_mat_y, mat_y, current_mat_z, mat_z, current_err, err
  identity = linalg_ops.eye(math_ops.cast(mat_a_size, dtypes.int32))
  mat_a = mat_a + ridge_epsilon * identity
  norm = math_ops.sqrt(math_ops.reduce_sum(mat_a * mat_a))
  mat_init_y = mat_a / norm
  mat_init_z = identity
  init_err = norm
  _, _, prev_mat_y, _, _, _, _ = control_flow_ops.while_loop(
      _iter_condition, _iter_body, [
          0, mat_init_y, mat_init_y, mat_init_z, mat_init_z, init_err,
          init_err + 1.0
      ])
  return prev_mat_y * math_ops.sqrt(norm)