Fix unbind op.

torch.unbind does not lower to aten.unbind so we don't need to map it.

The test passed with the previous implementation but I created a
slightly more efficient one anyway.

experimental/torch_xla2/test/test_ops.py

@@ -190,7 +190,6 @@
  "take_along_dim",
  "to_sparse", # We are not supporting sparse tensors yet.
  "triu",
- "unbind",
  "unfold_copy",
  "unfold",
  "unique_consecutive",

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -2396,13 +2396,9 @@ def _aten_trunc(a):
  return jnp.trunc(a)
 
 
-@op(torch.ops.aten.unbind)
 @op(torch.ops.aten.unbind_copy)
 def _aten_unbind(a, dim=0):
- return tuple(
- _aten_squeeze_dim(jax.lax.index_in_dim(a, i, axis=dim), dim)
- for i in range(a.shape[dim])
- )
+ return [jax.lax.index_in_dim(a, i, dim, keepdims=False) for i in range(a.shape[dim])]
 
 
 # NOTE: skip aten.upsample_nearest2d and aten.upsample_bilinear2d