Fixes pytorch/xla#7398

torchax/test/test_ops.py

@@ -13,7 +13,6 @@
     "_segment_reduce",
     "bincount",  # NOTE: dtype for int input torch gives float. This is weird.
     "byte",
-    "cat",
     "cholesky_solve",
     "diagonal_copy",
     "geqrf",

torchax/torchax/ops/jaten.py

@@ -1248,6 +1248,11 @@ def _aten_relu(self):
 
 @op(torch.ops.aten.cat)
 def _aten_cat(tensors, dims=0):
+  # handle it as a special case if the first tensor is empty.
+  # torch.cat will ignore the empty tensor, while jnp.concatenate
+  # will error if the dims > 0.
+  if len(tensors) > 0 and tensors[0].ndim == 1 and tensors[0].shape[0] == 0:
+    return jnp.concatenate(tensors[1:], dims)
   return jnp.concatenate(tensors, dims)