[Fori Loop] Inconsistent Shape Behavior about xla HOT 2 OPEN

@ManfeiBai can you take a look?

from xla.

Hi, @huzama, thanks for pointing this

yes, I have reproduced the same error with your instruction locally too:

• self.quantity = torch.tensor(0, device=self.device) failed
• self.quantity = torch.tensor([0], device=self.device) passed
• and noticed self.quantity = torch.tensor(5, device=self.device) passed

tried again with more try:

• self.quantity = torch.tensor(0, device=self.device) failed: https://gist.github.com/ManfeiBai/ea9cf0f0db409f4bcdd45f6af65c638d
• self.quantity = torch.tensor(1, device=self.device) failed: https://gist.github.com/ManfeiBai/cf084d713bc4871d8e5f95fe7913e364
• self.quantity = torch.tensor(2, device=self.device) passed: https://gist.github.com/ManfeiBai/09933c57f3afbbfcc0fdccf95b84ac8b
• self.quantity = torch.tensor(3, device=self.device) passed: https://gist.github.com/ManfeiBai/3b13d4ea1068e2b497e60ccb80b66a1b
• self.quantity = torch.tensor(5, device=self.device) passed: https://gist.github.com/ManfeiBai/dd711bf78acfc412b3ac4e10283076d1

based on above link's log's xlacomputation for body and cond, we notice that xlacomputation generation would missed/ignore tensor(0) and tensor(1), and try to create a new constant locally to use

for how to print xlacomputation(), not sure if anyone need this, but post it here and will add in while_loop document later:
adding these code after line1 and line2:

# after line1
  body_hlo_print = xb.get_computation_hlo(body_computation)
  print("body computation: !!!!!!!!!")
  print(body_hlo_print)


# after line2
  cond_hlo_print = xb.get_computation_hlo(cond_computation)
  print("cond computation: !!!!!!!!!")
  print(cond_hlo_print)

from xla.