Fix shaped broadcast

EnzymeAD · Dec 27, 2024 · f2094b3 · f2094b3
1 parent b2d055f
commit f2094b3
Showing 1 changed file with 12 additions and 17 deletions.
diff --git a/src/enzyme_ad/jax/Passes/ArithRaising.cpp b/src/enzyme_ad/jax/Passes/ArithRaising.cpp
@@ -96,24 +96,19 @@ struct ArithRaisingPass : public ArithRaisingPassBase<ArithRaisingPass> {
     op->walk([=](enzyme::BroadcastOp broadcastOp) {
       OpBuilder builder(broadcastOp);
       Value newBroadcastOp;
-      if (use_stablehlo) {
-        SmallVector<int64_t> broadcastDims;
-        auto shape =
-            broadcastOp.getInput().getType().cast<TensorType>().getShape();
-        broadcastDims.reserve(shape.size());
-        for (auto en : llvm::enumerate(shape)) {
-          // original dimensions end up one further because the batch dimension
-          // is prepended:
-          broadcastDims.push_back(en.index() + 1);
-        }
-        newBroadcastOp = builder.create<stablehlo::BroadcastInDimOp>(
-            broadcastOp.getLoc(), broadcastOp.getType(), broadcastOp.getInput(),
-            builder.getDenseI64ArrayAttr(broadcastDims));
-      } else {
-        newBroadcastOp = builder.create<mhlo::BroadcastOp>(
-            broadcastOp.getLoc(), broadcastOp.getInput(),
-            builder.getI64TensorAttr({broadcastOp.getWidth()}));
+      assert(use_stablehlo);
+      SmallVector<int64_t> broadcastDims;
+      auto shape =
+          broadcastOp.getInput().getType().cast<TensorType>().getShape();
+      broadcastDims.reserve(shape.size());
+      for (auto en : llvm::enumerate(shape)) {
+        // original dimensions end up one further because the batch dimension
+        // is prepended:
+        broadcastDims.push_back(en.index() + 1);
       }
+      newBroadcastOp = builder.create<stablehlo::BroadcastInDimOp>(
+          broadcastOp.getLoc(), broadcastOp.getType(), broadcastOp.getInput(),
+          builder.getDenseI64ArrayAttr(broadcastDims));
       broadcastOp.replaceAllUsesWith(newBroadcastOp);
       broadcastOp.erase();
     });