Fix and test reshape of constants

EnzymeAD · Mar 12, 2024 · 98f5dbd · 98f5dbd
1 parent 7833e8c
commit 98f5dbd
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -1600,7 +1600,7 @@ struct ReshapeSimplify : public OpRewritePattern<mlir::stablehlo::ReshapeOp> {
     if (inp) {
       rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
           op, op.getType(),
-          inp.isSplat() ? inp.resizeSplat(inp.getType())
+          inp.isSplat() ? inp.resizeSplat(op.getType())
                         : inp.reshape(op.getType()));
       return success();
     }

diff --git a/test/lit_tests/reshapeconst.mlir b/test/lit_tests/reshapeconst.mlir
@@ -0,0 +1,14 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+  func.func @main() -> tensor<1xf32> {
+    %concat = stablehlo.constant dense<3.140000e+00> : tensor<f32>
+    %conv = stablehlo.reshape %concat : (tensor<f32>) -> tensor<1xf32>
+    return %conv : tensor<1xf32>
+  }
+}
+
+// CHECK:  func.func @main() -> tensor<1xf32> {
+// CHECK-NEXT:    %0 = stablehlo.constant dense<3.140000e+00> : tensor<1xf32>
+// CHECK-NEXT:    return %0 : tensor<1xf32>
+// CHECK-NEXT:  }