fix a bug in tensor.flatten

Author: yaochengji

torchax/test/test_functions.py

@@ -60,6 +60,12 @@ def test_bernoulli_inplace(self):
       a = torch.randn((2,3))
       a.bernoulli_(0.4)
 
+  def test_flatten(self):
+    with self.env:
+      a = torch.randn((2,3,4))
+      a = a.flatten(0, 1)
+      self.assertEqual(tuple(a.shape), (6, 4))
+
   def test_rnn(self):
     model = SeqModel()
     x = torch.randn((2, 100, 20))

torchax/torchax/tensor.py

@@ -100,7 +100,7 @@ def flatten(self, start_dim=0, end_dim=-1):
     if end_dim == -1:
       end_dim = self.ndim
     new_shape = (
-        self._elem.shape[:start_dim] + (-1,) + self._elem.shape[end_dim:])
+        self._elem.shape[:start_dim] + (-1,) + self._elem.shape[end_dim + 1:])
     new_elem = jnp.reshape(self._elem, new_shape)
     return Tensor(new_elem, self._env)
     # return torch.reshape(self, new_shape)
@@ -371,6 +371,8 @@ def _handle_tensor_constructor(self, func, args, kwargs):
           return func(*args, **kwargs)
       with jax.default_device(jax_device):
         op = self._ops.get(func)
+        if op is None and isinstance(func, torch._ops.OpOverload):
+          op = self._ops.get(func.overloadpacket)
         res = op.func(*args, **kwargs)
         if isinstance(res, jax.Array):
           res = Tensor(res, self)