fix: update to new version

Author: avik-pal

src/enzyme_ad/jax/Utils.cpp

@@ -672,8 +672,6 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
     return perm.size() == 2 && perm[0] == 1 && perm[1] == 0;
   };
 
-  // TODO: check for dot_general as well
-
   if (auto broadcastOp = dyn_cast<stablehlo::BroadcastInDimOp>(op)) {
     auto operand = broadcastOp.getOperand();
     auto operandTy = cast<RankedTensorType>(operand.getType());
@@ -710,34 +708,33 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
     }
   }
 
-  // A(A^T) will always be symmetric
-  if (auto dotOp = dyn_cast_or_null<stablehlo::DotGeneralOp>(op)) {
+  // A x (A^T) / (A^T) x A will always be symmetric
+  if (auto dotOp = dyn_cast<stablehlo::DotGeneralOp>(op)) {
+    auto dotDimNumbers = dotOp.getDotDimensionNumbers();
     auto lhs = dotOp.getLhs();
     auto rhs = dotOp.getRhs();
 
-    auto lhsDims = dotOp.getDotDimensionNumbers().getLhsContractingDimensions();
-    auto rhsDims = dotOp.getDotDimensionNumbers().getRhsContractingDimensions();
+    auto lhsCDims = dotDimNumbers.getLhsContractingDimensions();
+    auto rhsCDims = dotDimNumbers.getRhsContractingDimensions();
 
-    if (auto lhsType = dyn_cast_or_null<ShapedType>(lhs.getType());
-        lhsType && lhsType.hasRank() && lhsType.getRank() == 2) {
+    if (dotDimNumbers.getLhsBatchingDimensions().size() == 0 &&
+        dotDimNumbers.getRhsBatchingDimensions().size() == 0 &&
+        lhsCDims.size() == 1 && rhsCDims.size() == 1) {
+      if (lhs == rhs && lhsCDims[0] == rhsCDims[0]) {
+        return State::GUARANTEED;
+      }
 
-      if (rhs == lhs) {
-        if (lhsDims.size() == 1 && lhsDims == rhsDims) {
+      if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
+        if (isTrueTranspose(lhsT) && lhsT.getOperand() == rhs &&
+            lhsCDims[0] == 1 - rhsCDims[0]) {
           return State::GUARANTEED;
         }
       }
 
-      if (lhsDims.size() == 1 && lhsDims != rhsDims) {
-        if (auto rhsT = rhs.getDefiningOp<stablehlo::TransposeOp>()) {
-          auto rhsInput = rhsT.getOperand();
-          if (rhsInput == lhs && isTrueTranspose(rhsT))
-            return State::GUARANTEED;
-        }
-
-        if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
-          auto lhsInput = lhsT.getOperand();
-          if (lhsInput == rhs && isTrueTranspose(lhsT))
-            return State::GUARANTEED;
+      if (auto rhsT = rhs.getDefiningOp<stablehlo::TransposeOp>()) {
+        if (isTrueTranspose(rhsT) && rhsT.getOperand() == lhs &&
+            lhsCDims[0] == 1 - rhsCDims[0]) {
+          return State::GUARANTEED;
         }
       }
     }

test/lit_tests/structured_tensors/transpose_symmetric.mlir

@@ -55,7 +55,7 @@ func.func @pass4(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 
 // CHECK: func.func @pass4(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 // CHECK-NEXT:   %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
-// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.guaranteed_symmetric = true} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT:   return %1 : tensor<2x2xf32>
 // CHECK-NEXT: }
 
@@ -66,7 +66,7 @@ func.func @pass5(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 }
 
 // CHECK: func.func @pass5(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
-// CHECK-NEXT:   %0 = stablehlo.dot_general %arg0, %arg0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.guaranteed_symmetric = true} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %0 = stablehlo.dot_general %arg0, %arg0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT:   return %0 : tensor<2x2xf32>
 // CHECK-NEXT: }
 
@@ -93,7 +93,7 @@ func.func @fail2(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 
 // CHECK: func.func @fail2(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 // CHECK-NEXT:   %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
-// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] {enzymexla.guaranteed_symmetric = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT:   %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT:   return %2 : tensor<2x2xf32>
 // CHECK-NEXT: }