%0 = "tf.GatherV2"(%arg0, %arg1, %arg2) {attr_map = "0:batch_dims", batch_dims = 0 : i64, device = ""} : (tensor<1024x3x4x3xf32>, tensor<1xi32>, tensor<i32>) -> tensor<1x3x4x3xf32>
    return %0 : tensor<1x3x4x3xf32>
  }

    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
    %1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>
    return %1: tensor<1x3xbf16>
  }

      attributes {tf_quant.quantized_ops = ${quantized_ops}}
  {
    %accum_out = "tf.GatherV2"(%weight, %input, %axis) {
      batch_dims = 0 : i64, attr_map = "batch_dims:0"} : (tensor<*xi8>, tensor<*xi32>, tensor<i32>) -> tensor<*xi8>

    %accum_out_new = "tf.Cast"(%accum_out) : (tensor<*xi8>) -> tensor<*xf32>

      attributes {tf_quant.quantized_ops = ["Gather"]} {

    %out = "tf.GatherV2"(%weight, %input, %axis) {
      batch_dims = 0 : i64, attr_map = "batch_dims:0"} : (tensor<*xi8>, tensor<*xi32>, tensor<i32>) -> tensor<*xi8>

    // Requantize as the output quantization params can be different than the input for Gather ops.

// CHECK:     %17 = "tfl.gather"(%arg13, %14) <{axis = 0 : i32, batch_dims = 0 : i32}> {tac.device = "DARWINN", tac.inference_type = "FLOAT"} : (tensor<5xi32>, tensor<?xi32>) -> tensor<?xi32>
// CHECK:     %18 = tfl.add %arg14, %17 {fused_activation_function = "NONE", tac.device = "DARWINN", tac.inference_type = "FLOAT"} : tensor<?xi32>

// Here, $params and $indices needs to be ranked so that $axis and $batch_dims
// attributes can be converted from TensorFlow axis format supporting negative
// indexing to the HLO format.
def LegalizeGatherV2 :
  Pat<(TF_GatherV2Op AnyRankedTensor:$params, AnyRankedTensor:$indices,
        (ConstantLikeMatcher ElementsAttr:$axis), $batch_dims),
      (MHLO_TorchIndexSelectOp $params, $indices,

  %3 = "tf.AddV2"(%arg1, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %4 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>

  %383 = "tf.GatherV2"(%11, %382#1, %14) {batch_dims = 0 : i64, device = ""} : (tensor<1xi64>, tensor<?xi64>, tensor<i32>) -> tensor<?xi64>
  %384 = "tf.GatherV2"(%12, %383, %14) {batch_dims = 0 : i64, device = ""} : (tensor<1xi64>, tensor<?xi64>, tensor<i32>) -> tensor<?xi64>

                           (TFL_GatherNdOp $params, $indices)>;

def LegalizeGatherV2 : Pat<
  (TF_GatherV2Op $params, $indices, (Arith_ConstantOp ElementsAttr:$axis), $batch_dims),
  (TFL_GatherOp $params, $indices, ExtractSingleElementAsInt32:$axis,
    (convertIntAttrTo32Bit $batch_dims))>;

def LegalizeFloorDiv : Pat<(TF_FloorDivOp $l, $r), (TFL_FloorDivOp $l, $r)>;

def LegalizeNotEqual : Pat<

LogicalResult GatherV2Op::verify() {
  GatherV2Op op = *this;
  int64_t batch_dims = op.getBatchDims();
  if (auto ty = mlir::dyn_cast<RankedTensorType>(op.getIndices().getType())) {
    int64_t rank = ty.getRank();
    if (batch_dims > rank || batch_dims < -rank)
      return op.emitOpError()
             << "batch_dims (" << batch_dims << ") must be in range [" << -rank
             << ", " << rank + 1 << ")";