}

  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),

std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op) {
  auto scale_spec = std::make_unique<OpQuantScaleSpec>();
  if (llvm::isa<mlir::stablehlo::BroadcastInDimOp,
                mlir::stablehlo::ConcatenateOp,
                mlir::stablehlo::DynamicReshapeOp,
                mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,
                mlir::stablehlo::PadOp, mlir::stablehlo::ReduceWindowOp,

class RewriteQuantizedConcatenateOp
    : public OpRewritePattern<stablehlo::ConcatenateOp> {
 public:
  using OpRewritePattern<stablehlo::ConcatenateOp>::OpRewritePattern;

  LogicalResult match(stablehlo::ConcatenateOp op) const override {
    return success(IsOpFullyQuantized(op));
  }

  void rewrite(stablehlo::ConcatenateOp op,
               PatternRewriter& rewriter) const override {

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;

    // offset. This also works if shift is not constant.
    // DynamicSliceOp requires the sizes being integer, and we can get the
    // information from input shape.
    auto concat = b.create<ConcatenateOp>(
        ValueRange{op.getInput(), op.getInput()}, b.getI64IntegerAttr(axis));
    Value zero = b.create<mhlo::ConstantOp>(
        b.getIntegerAttr(getElementTypeOrSelf(offset.getType()), 0));

          conv_op.getPrecisionConfigAttr());

      conv_results.push_back(conv_result);
    }

    auto final_output = rewriter.create<mhlo::ConcatenateOp>(
        conv_op.getLoc(), conv_results,
        rewriter.getI64IntegerAttr(dnums.getOutputFeatureDimension()));
    rewriter.replaceOp(conv_op, final_output.getResult());
    return mlir::success();
  }