pytorch

1
def while_test(a: Tensor,
2
    i: Tensor) -> Tensor:
3
  a0 = a
4
  i0 = i
5
  _0 = bool(torch.lt(i, 3))
6
  while _0:
7
    a1 = torch.mul_(a0, a0)
8
    i1 = torch.add_(i0, 1)
9
    _0, a0, i0 = bool(torch.lt(i1, 3)), a1, i1
10
  return a0
11