code: def sqrt(x):
      """
      EXPERIMENTAL Computes the upper triangular square root of an
      symmetric arrowhead matrix.
  
      :param SymmArrowhead x: an symmetric arrowhead matrix
      :return: the square root of `x`
      :rtype: TriuArrowhead
      """
      assert isinstance(x, SymmArrowhead)
      head_size = x.top.size(0)
      if head_size == 0:
          return TriuArrowhead(x.top, x.bottom_diag.sqrt())
  
      A, B = x.top[:, :head_size], x.top[:, head_size:]
      # NB: the complexity is O(N * head_size^2)
      # ref: https://en.wikipedia.org/wiki/Schur_complement#Background
      Dsqrt = x.bottom_diag.sqrt()
  
      # On cholesky error, retry with smaller tail part B.
      num_attempts = 6
      for i in range(num_attempts):
          B_Dsqrt = B / Dsqrt.unsqueeze(-2)  # shape: head_size x N
          schur_complement = A - B_Dsqrt.matmul(
              B_Dsqrt.t()
          )  # complexity: head_size^2 x N
          # we will decompose schur_complement to U @ U.T (so that the sqrt matrix
          # is upper triangular) using some `flip` operators:
          #   flip(cholesky(flip(schur_complement)))
          try:
              top_left = torch.flip(
                  torch.linalg.cholesky(
                      torch.flip(schur_complement, (-2, -1))
                  ),
                  (-2, -1),
              )
              break
          except RuntimeError:
              B = B / 2
              continue
          raise RuntimeError(
              "Singular schur complement in computing Cholesky of the input"
              " arrowhead matrix"
          )
  
      top_right = B_Dsqrt
      top = torch.cat([top_left, top_right], -1)
      bottom_diag = Dsqrt
      return TriuArrowhead(top, bottom_diag)
  
imports: from collections import namedtuple
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: sqrt(x) and returns: TriuArrowhead, TriuArrowhead