fmt

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3790,33 +3790,33 @@ struct BroadcastToReshape final
     else {
       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 {
+      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();
   }
 };