def _apply_ith_rotation(i, args):
  """ Applies the two-dimensional ("Givens") rotation matrix specified by
  cs[i] and sn[i] (the cosine and sine factors of the rotation) to the
  two-vector specified by the ith and i+1th components of h. Returns
  the rotated vector along with the unmodified cs and sn, which is
  necessary for Jax typing reasons.
  """
  h, cs, sn = args
  x1 = h[i]
  y1 = h[i + 1]
  x2 = cs[i].conj() * x1 - sn[i].conj() * y1
  y2 = sn[i] * x1 + cs[i] * y1
  h = h.at[i].set(x2)
  h = h.at[i + 1].set(y2)
  return h, cs, sn