test range

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

@@ -81,11 +81,13 @@ def test_fori_loop_tpu_addition(self):
     upper = torch.tensor([30], dtype=torch.int32, device=device)
     one_value = torch.tensor([1], dtype=torch.int32, device=device)
     init_val_list = (init_val, one_value)
+    lowers = torch.tensor([1, 1, 1], dtype=torch.int32, device=device) # lower, init_val, one_value
 
     def body_fun(a, b):
       return torch.add(a, b) # [0])
     # _, _, res, _ = fori_loop(lower, upper, body_fun, init_val, one_value) # init_val_list) # init_val)
-    A, B, res, D = fori_loop(lower, upper, body_fun, init_val, one_value) # init_val_list) # init_val)
+    # A, B, res, D = fori_loop(lower, upper, body_fun, init_val, one_value) # init_val_list) # init_val)
+    A, B, res, D = fori_loop(upper, lowers) # lower, upper, body_fun, init_val, one_value)
     print("result: ", res) # init_val_
     print("A: ", A) # lower_
     print("B: ", B) # upper_

torch_xla/experimental/fori_loop.py

@@ -12,7 +12,7 @@
 from torch._higher_order_ops.while_loop import while_loop_op
 
 
-def fori_loop(lower, upper, body_fun, *init_vals): # *init_val):
+def fori_loop(lower, lowers):#  upper, body_fun, *init_vals): # *init_val):
 
   device = xm.xla_device()
   # limit_value = upper
@@ -52,7 +52,7 @@ def body_fn(upper, lowers): # , *init_vals):
     # body_fun(one_value_original, init_val)) # body_fun(lower, init_val))
 
   # res = while_loop(cond_fn, body_fn, (upper, lower, *init_vals))
-  lowers = (lower, *init_vals)
+  # lowers = (lower, *init_vals)
   res = _xla_while_loop(cond_fn, body_fn, (upper, lowers)) # , *init_vals))
   # print("upper: ", upper)
   # print("lower: ", lower)