shape_of_op.replaceAllUsesWith(new_shape_of);
      rewriter.replaceOp(mul_op, {new_conv});
    }

    return success();
  }
};

class FuseMhloConvolutionPass
    : public PassWrapper<FuseMhloConvolutionPass, OperationPass<func::FuncOp>> {
 public:
  StringRef getArgument() const final { return "fuse-mhlo-convolution-pass"; }
  StringRef getDescription() const final {