def stable_pos_enc(x, n):
  """A stable posenc for very high degrees, courtesy of Sameer Agrawal."""
  sin_x = np.sin(x)
  cos_x = np.cos(x)
  output = []
  rotmat = np.array([[cos_x, -sin_x], [sin_x, cos_x]], dtype='double')
  for _ in range(n):
    output.append(rotmat[::-1, 0, :])
    rotmat = np.einsum('ijn,jkn->ikn', rotmat, rotmat)
  return np.reshape(np.transpose(np.stack(output, 0), [2, 1, 0]), [-1, 2 * n])
A stable posenc for very high degrees, courtesy of Sameer Agrawal.
internal/mip_test.py