add unbind (#10)

* add unbind

* add tests and implements unbind

* fix x.unbind()

---------

Co-authored-by: matt hyatt <mhyatt000>

Author: mhyatt000

tests/test_core.py

@@ -77,6 +77,13 @@ def test_zeros_like():
   t2j_function_test(lambda x: torch.zeros_like(x), [(2, 3)])
 
 
+def test_unbind():
+  t2j_function_test(lambda x: torch.unbind(x)[0], [(2, 3)])
+  t2j_function_test(lambda x: torch.unbind(x, dim=1)[1], [(2, 3)])
+  t2j_function_test(lambda x: x.unbind()[0], [(2, 3)])
+  t2j_function_test(lambda x: x.unbind(1)[1], [(2, 3)])
+
+
 def test_oneliners():
   t2j_function_test(lambda x: torch.pow(x, 2), [()])
   t2j_function_test(lambda x: torch.pow(x, 2), [(3,)])

torch2jax/__init__.py

@@ -159,6 +159,7 @@ def sum(*args, **kwargs): return torch.sum(*args, **kwargs)
   def transpose(*args, **kwargs): return torch.transpose(*args, **kwargs)
   def view(self, *shape): return Torchish(jnp.reshape(self.value, shape))
   reshape = view
+  def unbind(*args, **kwargs): return torch.unbind(*args, **kwargs)
   # fmt: on
 
   def add_(self, other):
@@ -489,6 +490,11 @@ def tensor(data, dtype=None, device=None, requires_grad=False, pin_memory=False)
   )
 
 
+@implements(torch.unbind, Torchishify_output=False)
+def unbind(input, dim=0) -> Sequence[Torchish]:
+  return tuple(Torchish(input.value[(slice(None),) * dim + (i,)]) for i in range(input.value.shape[dim]))
+
+
 @implements(torch.zeros)
 def zeros(*args, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False):
   assert out is None, "TODO: implement out"