Fix the Op info test of cat

Author: qianminj123

torchax/test/test_ops.py

@@ -14,7 +14,6 @@
     "_segment_reduce",
     "bincount", # NOTE: dtype for int input torch gives float. This is weird.
     "byte",
-    "cat",
     "cholesky_solve",
     "diagonal_copy",
     "geqrf",
@@ -204,6 +203,10 @@ def test_reference_eager(self, device, dtype, op):
 
       # print("[DEBUG] sample_input: ", sample_input)
 
+      if op.name == "cat":
+        dim = sample_input.kwargs.get('dim')
+        if dim and dim >= sample_input.input[0].dim():
+            continue
       # TODO: this is a workaround to skip int64 cast for linspace
       # reference: https://github.com/pytorch/xla/issues/7505#issuecomment-2400895692 and subsequent comments
       # we have opened a bug in pytorch: https://github.com/pytorch/pytorch/issues/137546