return success();
  }
};

class SanitizeGatherOpOutputToI4 : public OpRewritePattern<TFL::GatherOp> {
 public:
  using OpRewritePattern<TFL::GatherOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::GatherOp op,
                                PatternRewriter &rewriter) const override {
    auto const_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());

             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp, TF::ExpandDimsOp,
             TF::SqueezeOp, TF::TransposeOp>(op);
}

bool IsOpWithQuantizableTrait(Operation* op) {
  // Supported quantizable ops.
  return isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::MatMulOp, TF::Conv2DOp,
             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp,

                mlir::stablehlo::PadOp, mlir::stablehlo::ReduceWindowOp,
                mlir::stablehlo::ReshapeOp, mlir::stablehlo::SelectOp,
                mlir::stablehlo::SliceOp, mlir::stablehlo::TransposeOp>(op)) {
    scale_spec->has_same_scale_requirement = true;
  }
  if (llvm::isa<mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,

      auto batch_axes_tensor =
          builder.create<TFL::ConstOp>(builder.getI64TensorAttr(
              dot_dimensions_info.batch_dimensions().AxesArray()));
      auto batch_dims = builder.create<TFL::GatherOp>(
          RankedTensorType::get(
              {static_cast<int>(
                  dot_dimensions_info.batch_dimensions().AxesArray().size())},
              builder.getIntegerType(32)),

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/test.h"

namespace mlir::quant::stablehlo {
namespace {

using ::mlir::stablehlo::GatherOp;
using ::testing::IsEmpty;
using ::testing::IsTrue;
using ::testing::NotNull;
using ::testing::Pair;
using ::testing::UnorderedElementsAre;

using IsOpQuantizableStableHloTest = ::mlir::quant::QuantizationTestBase;

  absl::flat_hash_set<int> quantizable_operands;
  if (isa<TF::DepthwiseConv2dNativeOp, TF::Conv2DOp, TF::Conv3DOp, TF::MatMulOp,
          TF::BatchMatMulOp>(op)) {
    quantizable_operands.insert(1);
  } else if (isa<TF::GatherOp>(op)) {
    quantizable_operands.insert(0);
  } else if (auto einsum_op = dyn_cast<TF::EinsumOp>(op)) {
    if (IsEinsumSupportedByXlaDotV2(einsum_op.getEquationAttr())) {
      quantizable_operands.insert(1);
    }

using ::mlir::stablehlo::ConcatenateOp;
using ::mlir::stablehlo::ConvolutionOp;
using ::mlir::stablehlo::DotGeneralOp;
using ::mlir::stablehlo::DynamicBroadcastInDimOp;
using ::mlir::stablehlo::GatherOp;
using ::mlir::stablehlo::GetDimensionSizeOp;
using ::mlir::stablehlo::ReshapeOp;
using ::mlir::stablehlo::UniformQuantizeOp;
using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::QuantizedDimension;

  auto gather_lhs = rewriter->create<TF::GatherOp>(
      loc,
      RankedTensorType::get({static_cast<int>(bl_index.size())},
                            rewriter->getIntegerType(32)),
      shape_lhs->getResults()[0], lhs_index_tensor->getResults()[0], true_attr);
  auto rhs_index_tensor = createI32ConstantOp(r_index, loc, rewriter);
  auto gather_rhs = rewriter->create<TF::GatherOp>(
      loc,

      auto batch_axes_tensor =
          builder.create<TF::ConstOp>(builder.getI64TensorAttr(
              dot_dimensions_info.batch_dimensions().AxesArray()));
      auto batch_dims = builder.create<TF::GatherOp>(
          RankedTensorType::get(
              {static_cast<int>(
                  dot_dimensions_info.batch_dimensions().AxesArray().size())},
              builder.getIntegerType(32)),

// Condition 4 - shape of slice should be same with shape of input on the
// offset dimensions.
class RewriteQuantizedGatherOp : public OpRewritePattern<stablehlo::GatherOp> {
 public:
  using OpRewritePattern<stablehlo::GatherOp>::OpRewritePattern;

  LogicalResult match(stablehlo::GatherOp op) const override {
    const Type input_type = op.getOperand().getType();
    const Type output_type = op.getResult().getType();