einsum: conditionally do squeeze before transpose (#4079)

Author: lakhinderwalia

src/onnx/parse_einsum.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -619,7 +619,8 @@ struct parse_einsum : op_parser<parse_einsum>
                 perm.push_back(row_output[i]);
         }
 
-        op = info.add_instruction(make_op("squeeze", {{"axes", sq_axes}}), op);
+        if(not sq_axes.empty())
+            op = info.add_instruction(make_op("squeeze", {{"axes", sq_axes}}), op);
 
         if(not perm.empty())
         {

test/onnx/einsum_permute_sd3_test.onnx

@@ -0,0 +1,21 @@
+
+einsum_permute_sd3_test:�

+-
+
x
y"Einsum*
+equation"nhwpqc->nchpwq�
einsum_permute_sd3_testZ#
+
x
+

+

+@
+@
+
+
+b#
+
y
+

+

+
+@
+
+@
+B

test/onnx/gen_onnx.py

@@ -2614,6 +2614,20 @@ def einsum_permute_test():
     return ([node], [x], [y])
 
 
+def einsum_permute_sd3_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT,
+                                      [1, 64, 64, 2, 2, 16])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT,
+                                      [1, 16, 64, 2, 64, 2])
+
+    node = onnx.helper.make_node('Einsum',
+                                 inputs=['x'],
+                                 outputs=['y'],
+                                 equation='nhwpqc->nchpwq')
+
+    return ([node], [x], [y])
+
+
 @onnx_test()
 def einsum_summation_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 3])