code: def triu_gram(x):
      """
      EXPERIMENTAL Computes the gram matrix `x.T @ x` from an upper-triangular
      arrowhead matrix `x`.
  
      :param TriuArrowhead x: an upper-triangular arrowhead matrix.
      :return: the square of `x`
      :rtype: TriuArrowhead
      """
      assert isinstance(x, TriuArrowhead)
      head_size = x.top.size(0)
      if head_size == 0:
          return x.bottom_diag.pow(2)
  
      A, B = x.top[:, :head_size], x.top[:, head_size:]
      top = A.t().matmul(x.top)
      bottom_left = top[:, head_size:].t()
      # the following matmul operator is O(N^2 x head_size)
      bottom_right = B.t().matmul(B) + x.bottom_diag.pow(2).diag()
      return torch.cat(
          [top, torch.cat([bottom_left, bottom_right], -1)], 0
      )
  
imports: from collections import namedtuple
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: triu_gram(x) and returns: torch, x