Add extra APIs. (pytorch#6586)

experimental/torch_xla2/test/test_extra.py

@@ -0,0 +1,31 @@
+import unittest
+import torch
+import torch.nn.functional as F
+import jax
+import torch_xla2
+from torch_xla2 import tensor, extra
+
+
+class ExtraTest(unittest.TestCase):
+
+  def setUp(self):
+    torch.manual_seed(0)
+
+  def test_fori_loop(self):
+    a = tensor.move_to_device(torch.ones((10, 10)))
+
+    def body(i, c):
+        return c + a[i]
+
+    init_val = tensor.move_to_device(torch.zeros(10))
+    res = extra.fori_loop(0, 10, body, init_val)
+
+    expect = torch.ones(10) * 10
+
+    self.assertTrue(torch.allclose(tensor.j2t(res._elem), expect))
+
+
+
+
+if __name__ == '__main__':
+  unittest.main()

experimental/torch_xla2/torch_xla2/extra.py

@@ -0,0 +1,41 @@
+import jax
+import functools
+from torch.utils import _pytree as pytree
+from torch_xla2 import tensor
+
+
+def call_jax(jax_function, *args, **kwargs):
+    # args, kwargs are torch tensors
+    # return val is torch tensor
+    args, kwargs = tensor.unwrap((args, kwargs))
+    res = jax_function(*args, **kwargs)
+    return tensor.wrap(res)
+
+
+def call_torch(torch_function, *args, **kwargs):
+    # args, kwargs are torch tensors
+    # return val is torch tensor
+    args, kwargs = tensor.wrap((args, kwargs))
+    res = torch_function(*args, **kwargs)
+    return tensor.unwrap(res)
+
+
+def fori_loop(lower, upper, body_fn, init_val, *, unroll=None):
+    """Torch fori_loop mimicking jax behavior.
+
+    Args:
+        lower: lower bound
+        upper: upperbound
+        init_val: init value (tree of torch.Tensors)
+        body_fn is a function that takes (int, a) -> a
+            where a is a pytree with torch.Tensors
+        unroll = False | True | int
+    """
+    jax_body = functools.partial(call_torch, body_fn)
+    return call_jax(
+        jax.lax.fori_loop, 
+        lower, 
+        upper, 
+        jax_body, 
+        init_val, 
+        unroll=unroll)