Also broadcast2reshape

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3784,9 +3784,33 @@ struct BroadcastToReshape final
         return failure();
     }
 
-    // replace with reshape
-    rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
-                                                      op.getOperand());
+    auto NT = op.getType();
+    stablehlo::ReshapeOp reshaped = nullptr;
+    for (auto u : op.getOperand().getUsers()) {
+      auto re = dyn_cast<stablehlo::ReshapeOp>(u);
+      if (!re)
+        continue;
+      if (re.getType() != NT)
+        continue;
+      reshaped = re;
+      break;
+    }
+    if (!reshaped) {
+      rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
+                                                        op.getOperand());
+    } else {
+      if (reshaped->getBlock() == op->getBlock()) {
+        if (op->isBeforeInBlock(reshaped)) {
+          rewriter.modifyOpInPlace(reshaped,
+                                   [&]() { reshaped->moveBefore(op); });
+        }
+        rewriter.replaceOp(op, reshaped);
+      } else {
+        rewriter.replaceOpWithNewOp<stablehlo::ReshapeOp>(op, op.getType(),
+                                                          op.getOperand());
+      }
+    }
+
     return success();
   }
 };