Repository: EnzymeAD/Enzyme-JAX (language: python)

Issue:
Missing differentiation rules for `einsum`, `unary_einsum`
I'm getting the following error in [EnzymeAD/Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) when trying to differentiate `stablehlo.einsum`:

```
julia> f = Reactant.compile(grad, (a′, b′))
error: could not compute the adjoint for this operation %2 = "stablehlo.einsum"(%1, %0) <{einsum_config = "ij,jk->ik"}> : (tensor<2x2xf64>, tensor<2x2xf64>) -> tensor<2x2xf64>
Pipeline failed
```

I open the issue here because I believe here is where the EnzymeMLIR rules are declared for HLO dialects right?


Discussion / hints:

Fix patch:
diff --git a/src/enzyme_ad/jax/Implementations/HLODerivatives.td b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
index 89b79a93..ae242884 100644
--- a/src/enzyme_ad/jax/Implementations/HLODerivatives.td
+++ b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
@@ -516,6 +516,30 @@ def GradDataConvBatchGroupCount : GlobalExpr</*needsprimal*/0, /*needsshadow*/0,
 
 def LT : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "ComparisonDirection::LT">;
 
+def EinsumConfig : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
+  op.getEinsumConfig();
+}]>;
+
+def GradUnaryEinsumConfig : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
+  StringRef config = op.getEinsumConfig();
+  auto [input, output] = config.split(StringRef("->"));
+  builder.getStringAttr(output + "->" + input);
+}]>;
+
+def GradEinsumConfigLhs : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
+  StringRef config = op.getEinsumConfig();
+  auto [input, output] = config.split(StringRef("->"));
+  auto [lhs, rhs] = input.split(',');
+  builder.getStringAttr(output + "," + rhs + "->" + lhs);
+}]>;
+
+def GradEinsumConfigRhs : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
+  StringRef config = op.getEinsumConfig();
+  auto [input, output] = config.split(StringRef("->"));
+  auto [lhs, rhs] = input.split(',');
+  builder.getStringAttr(output + "," + lhs + "->" + rhs);
+}]>;
+
 // Derivative rules
 def : HLODerivative<"AddOp", (Op $x, $y),
                     [
@@ -558,6 +582,15 @@ def : HLODerivative<"DotGeneralOp", (Op $lhs, $rhs),
                     (Add (SelectIfActive $lhs, (DotGeneral (ResultTypes), (Shadow $lhs), $rhs, (ResultDotDim), (ResultDotPrec)), (HLOConstantFP<"0">)), (SelectIfActive $rhs, (DotGeneral (ResultTypes), $lhs, (Shadow $rhs), (ResultDotDim), (ResultDotPrec)), (HLOConstantFP<"0">)))
                   >;
 
+def : HLODerivative<"EinsumOp", (Op $lhs, $rhs),
+                    [
+                        // TODO add support for complex numbers by conjugating `$rhs` and `$lhs` (only if complex tensors)
+                        (Einsum (TypeOf $lhs), (DiffeRet), $rhs, (GradEinsumConfigLhs)),
+                        (Einsum (TypeOf $rhs), (DiffeRet), $lhs, (GradEinsumConfigRhs)),
+                    ],
+                    (Add (SelectIfActive $lhs, (Einsum (ResultTypes), (Shadow $lhs), $rhs, (EinsumConfig)), (HLOConstantFP<"0">)), (SelectIfActive $rhs, (Einsum (ResultTypes), $lhs, (Shadow $rhs), (EinsumConfig)), (HLOConstantFP<"0">)))
+                  >;
+
 def : HLODerivative<"ExpOp", (Op $x), [(CheckedMul (DiffeRet), (Exp $x))]>;
 
 def : HLODerivative<"LogOp", (Op $x), [(CheckedDiv (DiffeRet), $x)]>;
@@ -641,3 +674,11 @@ def : HLODerivative<"TransposeOp", (Op $x),
                     ],
                     (SelectIfActive $x, (Transpose (ResultTypes), (Shadow $x), (Permutation)), (HLOConstantFP<"0">))
                   >;
+
+def : HLODerivative<"UnaryEinsumOp", (Op $x),
+                    [
+                        // TODO maybe replace `UnaryEinsum` for `Reshape` + `BroadcastInDim`
+                        (UnaryEinsum (TypeOf $x), (DiffeRet), (GradUnaryEinsumConfig)),
+                    ],
+                    (UnaryEinsum (ResultTypes), (Shadow $x), (EinsumConfig))
+                  >;


Test patch:
diff --git a/test/lit_tests/einsum.mlir b/test/lit_tests/einsum.mlir
new file mode 100644
index 00000000..4096c4fb
--- /dev/null
+++ b/test/lit_tests/einsum.mlir
@@ -0,0 +1,29 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func @main(%a : tensor<2x3xf32>, %b : tensor<4x3x5xf32>) -> tensor<4x2x5xf32> {
+    %c = "stablehlo.einsum"(%a,%b) {einsum_config = "ab,cbd->cad"} : (tensor<2x3xf32>, tensor<4x3x5xf32>) -> tensor<4x2x5xf32>
+    func.return %c : tensor<4x2x5xf32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<4x3x5xf32>, %arg3: tensor<4x3x5xf32>) -> (tensor<4x2x5xf32>, tensor<4x2x5xf32>) {
+// FORWARD-NEXT:    %0 = stablehlo.einsum %arg1, %arg2, config = "ab,cbd->cad" : (tensor<2x3xf32>, tensor<4x3x5xf32>) -> tensor<4x2x5xf32>
+// FORWARD-NEXT:    %1 = stablehlo.einsum %arg0, %arg3, config = "ab,cbd->cad" : (tensor<2x3xf32>, tensor<4x3x5xf32>) -> tensor<4x2x5xf32>
+// FORWARD-NEXT:    %2 = stablehlo.add %0, %1 : tensor<4x2x5xf32>
+// FORWARD-NEXT:    %3 = stablehlo.einsum %arg0, %arg2, config = "ab,cbd->cad" : (tensor<2x3xf32>, tensor<4x3x5xf32>) -> tensor<4x2x5xf32>
+// FORWARD-NEXT:    return %3, %2 : tensor<4x2x5xf32>, tensor<4x2x5xf32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<4x3x5xf32>, %arg2: tensor<4x2x5xf32>) -> (tensor<2x3xf32>, tensor<4x3x5xf32>) {
+// REVERSE-NEXT:    %cst = arith.constant dense<0.000000e+00> : tensor<4x2x5xf32>
+// REVERSE-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<2x3xf32>
+// REVERSE-NEXT:    %cst_1 = arith.constant dense<0.000000e+00> : tensor<4x3x5xf32>
+// REVERSE-NEXT:    %0 = arith.addf %arg2, %cst : tensor<4x2x5xf32>
+// REVERSE-NEXT:    %1 = stablehlo.einsum %0, %arg1, config = "cad,cbd->ab" : (tensor<4x2x5xf32>, tensor<4x3x5xf32>) -> tensor<2x3xf32>
+// REVERSE-NEXT:    %2 = arith.addf %1, %cst_0 : tensor<2x3xf32>
+// REVERSE-NEXT:    %3 = stablehlo.einsum %0, %arg0, config = "cad,ab->cbd" : (tensor<4x2x5xf32>, tensor<2x3xf32>) -> tensor<4x3x5xf32>
+// REVERSE-NEXT:    %4 = arith.addf %3, %cst_1 : tensor<4x3x5xf32>
+// REVERSE-NEXT:    return %2, %4 : tensor<2x3xf32>, tensor<4x3x5xf32>
+// REVERSE-NEXT:  }
diff --git a/test/lit_tests/einsumcomplex.mlir b/test/lit_tests/einsumcomplex.mlir
new file mode 100644
index 00000000..39278e99
--- /dev/null
+++ b/test/lit_tests/einsumcomplex.mlir
@@ -0,0 +1,28 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+// XFAIL: *
+
+module {
+  func.func @main(%a : tensor<2x3xcomplex<f32>>, %b : tensor<4x3x5xcomplex<f32>>) -> tensor<4x2x5xcomplex<f32>> {
+    %c = "stablehlo.einsum"(%a,%b) {einsum_config = "ab,cbd->cad"} : (tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) -> tensor<4x2x5xcomplex<f32>>
+    func.return %c : tensor<4x2x5xcomplex<f32>>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<2x3xcomplex<f32>>, %arg1: tensor<2x3xcomplex<f32>>, %arg2: tensor<4x3x5xcomplex<f32>>, %arg3: tensor<4x3x5xcomplex<f32>>) -> (tensor<4x2x5xcomplex<f32>>, tensor<4x2x5xcomplex<f32>>) {
+// FORWARD-NEXT:    %0 = complex.conj %arg0 : tensor<2x3xcomplex<f32>>
+// FORWARD-NEXT:    %1 = complex.conj %arg2 : tensor<4x3x5xcomplex<f32>>
+// FORWARD-NEXT:    %2 = stablehlo.einsum %arg1, %1, config = "ab,cbd->cad" : (tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) -> tensor<4x2x5xcomplex<f32>>
+// FORWARD-NEXT:    %3 = stablehlo.einsum %0, %arg3, config = "ab,cbd->cad" : (tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) -> tensor<4x2x5xcomplex<f32>>
+// FORWARD-NEXT:    %4 = stablehlo.add %2, %3 : tensor<4x2x5xcomplex<f32>>
+// FORWARD-NEXT:    %3 = stablehlo.einsum %arg0, %arg2, config = "ab,cbd->cad" : (tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) -> tensor<4x2x5xcomplex<f32>>
+// FORWARD-NEXT:    return %3, %4 : tensor<4x2x5xcomplex<f32>>, tensor<4x2x5xcomplex<f32>>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<2x3xcomplex<f32>>, %arg1: tensor<4x3x5xcomplex<f32>>, %arg2: tensor<4x2x5xcomplex<f32>>) -> (tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) {
+// REVERSE-NEXT:    %0 = complex.conj %arg0 : tensor<2x3xcomplex<f32>>
+// REVERSE-NEXT:    %1 = complex.conj %arg1 : tensor<4x2x5xcomplex<f32>>
+// REVERSE-NEXT:    %2 = stablehlo.einsum %arg2, %0, config = "cad,ab->cbd" : (tensor<4x2x5xcomplex<f32>>, tensor<2x3xcomplex<f32>>) -> tensor<4x3x5xcomplex<f32>>
+// REVERSE-NEXT:    %3 = stablehlo.einsum %arg2, %1, config = "cad,cbd->ab" : (tensor<4x2x5xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>) -> tensor<2x3xcomplex<f32>>
+// REVERSE-NEXT:    return %2, %3 : tensor<2x3xcomplex<f32>>, tensor<4x3x5xcomplex<f32>>
+// REVERSE-NEXT:  }
diff --git a/test/lit_tests/unaryeinsum.mlir b/test/lit_tests/unaryeinsum.mlir
new file mode 100644
index 00000000..42ecf26d
--- /dev/null
+++ b/test/lit_tests/unaryeinsum.mlir
@@ -0,0 +1,24 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func @main(%x : tensor<2x3xf32>) -> tensor<2xf32> {
+    %y = "stablehlo.unary_einsum"(%x) {einsum_config = "ab->a"} : (tensor<2x3xf32>) -> tensor<2xf32>
+    func.return %y : tensor<2xf32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
+// FORWARD-NEXT:    %0 = stablehlo.unary_einsum %arg1, config = "ab->a" : (tensor<2x3xf32>) -> tensor<2xf32>
+// FORWARD-NEXT:    %1 = stablehlo.unary_einsum %arg0, config = "ab->a" : (tensor<2x3xf32>) -> tensor<2xf32>
+// FORWARD-NEXT:    return %1, %0 : tensor<2xf32>, tensor<2xf32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<2xf32>) -> tensor<2x3xf32> {
+// REVERSE-NEXT:    %cst = arith.constant dense<0.000000e+00> : tensor<2xf32>
+// REVERSE-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<2x3xf32>
+// REVERSE-NEXT:    %0 = arith.addf %arg1, %cst : tensor<2xf32>
+// REVERSE-NEXT:    %1 = stablehlo.unary_einsum %0, config = "a->ab" : (tensor<2xf32>) -> tensor<2x3xf32>
+// REVERSE-NEXT:    %2 = arith.addf %1, %cst_0 : tensor<2x3xf32>
+// REVERSE-NEXT:    return %2 : tensor<2x3xf32>
+// REVERSE-NEXT:  }