[torchax] fix einsum op conversion

Author: yaochengji

torchax/test/test_ops.py

@@ -227,6 +227,15 @@ def test_reference_eager(self, device, dtype, op):
       run_export_and_compare(self, op, sample_input, check_output, 
                              ignore_indices=ignore_index)
 
+  def test_einsum(self):
+    args = (
+        "bshd,bthd->bsht",
+        torch.ones((1, 2, 4, 8), dtype=torch.float32),
+        torch.ones((1, 2, 4, 8), dtype=torch.float32),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.einsum, args, kwargs, check_dtype=True)
+
 
 instantiate_device_type_tests(TestOpInfo, globals(), only_for={'cpu'})
 

torchax/torchax/ops/jtorch.py

@@ -86,14 +86,14 @@ def _diag(input, diagonal=0):
 def _einsum(equation, *operands):
   def get_params(*a):
     inner_list = a[0]
-    if len(inner_list) == 1:
+    if not isinstance(inner_list, jax.Array):
+      if len(inner_list) == 1:
         A = inner_list
         return A
-    elif len(inner_list) == 2:
+      elif len(inner_list) == 2:
         A, B = inner_list
         return A, B
-    else:
-        return operands
+    return operands
   assert isinstance(equation, str), 'Only accept str equation'
   filtered_operands = get_params(*operands)
   return jnp.einsum(equation, *filtered_operands)