def rotate90(tensor, times):
  """Rotate tensor clockwise by 90 degrees a number of times.
  Assumes spatial dimensions are the last ones."""
  if times == 1:  # 90 deg
    return tensor.transpose(-2, -1).flip(-1)
  elif times == 2:  # 180 deg
    return tensor.flip(-2).flip(-1)
  elif times == 3:  # 270 deg
    return tensor.transpose(-2, -1).flip(-2)
  else:  # 0 deg, no change
    assert times == 0
    return tensor