std::unique_ptr<Pass> createFoldBroadcastPass();

// Creates a pass which fuses MHLO binary element-wise ops and convolution op.
std::unique_ptr<Pass> createFuseConvolutionPass();

// Creates a pass which applies various optimizations on MHLO IR.
std::unique_ptr<Pass> createOptimizePass();

// Creates a pass that finds quantization patterns and compose them to uniform

                               const bool add_fold_broadcast_pass) {
  // Rewrites some patterns for better performance.
  pm.addNestedPass<func::FuncOp>(createUnfuseBatchNormPass());
  pm.addNestedPass<func::FuncOp>(createFuseConvolutionPass());
  pm.addNestedPass<func::FuncOp>(createOptimizePass());
  // Conditionally enable below pass because this causes unfused convolutions
  // described in b/293149194. This problem is not replicated in

    if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                            std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createFuseConvolutionPass() {
  return std::make_unique<FuseMhloConvolutionPass>();
}

static PassRegistration<FuseMhloConvolutionPass> pass;

}  // namespace odml

  // Unfuse mhlo BatchNorm to primitive ops.
  pm.addNestedPass<mlir::func::FuncOp>(mlir::odml::createUnfuseBatchNormPass());
  // Fuse Conv + Mul to Conv.
  pm.addNestedPass<mlir::func::FuncOp>(mlir::odml::createFuseConvolutionPass());
  // Fold broadcast_in_dim + Mul.
  pm.addNestedPass<mlir::func::FuncOp>(mlir::odml::createFoldBroadcastPass());
  pm.addNestedPass<mlir::func::FuncOp>(