}

  auto lhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(lhs_new_concat_shape),
      all_dot_lhs, rewriter.getI64IntegerAttr(lhs_batch_dim));
  auto rhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(rhs_new_concat_shape),

namespace mlir::quant::stablehlo {

namespace {

using ::mlir::quant::FindUserOfType;
using ::mlir::quant::TryCast;
using ::mlir::stablehlo::AddOp;
using ::mlir::stablehlo::BroadcastInDimOp;
using ::mlir::stablehlo::ConcatenateOp;
using ::mlir::stablehlo::ConvolutionOp;
using ::mlir::stablehlo::DotGeneralOp;
using ::mlir::stablehlo::DynamicBroadcastInDimOp;
using ::mlir::stablehlo::GatherOp;
using ::mlir::stablehlo::GetDimensionSizeOp;