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;

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,