int batch_dims;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(op.node()->attrs(), "batch_dims", &batch_dims));
  if (batch_dims < 0) {
    // TODO(bdodson): Figure out if we can find the param rank here, like the
    // python implementation does.
    return errors::InvalidArgument(
        "C++ GatherV2 gradient does not support negative batch_dims.");
  }

  %0 = "tfl.dequantize"(%arg0) : (tensor<1000x1000x!quant.uniform<i8:f32, 7.812500e-03>>) -> tensor<1000x1000xf32>
  %1 = "tfl.gather"(%0, %arg1) {axis = 0 : i32, batch_dims = 0 : i32}: (tensor<1000x1000xf32>, tensor<1x1xi32>) -> tensor<1x1x1000xf32>
  func.return %1 : tensor<1x1x1000xf32>

// CHECK-NEXT: tfl.gather
// CHECK-NEXT: tfl.dequantize
}

      auto batch_dims = builder.create<TFL::GatherOp>(
          RankedTensorType::get(
              {static_cast<int>(
                  dot_dimensions_info.batch_dimensions().AxesArray().size())},
              builder.getIntegerType(32)),
          operand_shape, batch_axes_tensor, /*axis*/ 0, /*batch_dims*/ 0);
      flattend_shape_values.push_back(batch_dims);
    } else {

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">
    (ArgumentList $params, $indices, $axis),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"batch_dims"> $batch_dims))),
  [(IsNotInLiftedFunc $res), (IsConstTensor $params)], [], (addBenefit 1)>;

def LiftConv3D : Pat<

      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>

func.func @gather(%indices: tensor<?xi32>,
             %resource: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
  %0 = "tf.ResourceGather"(%resource, %indices) {batch_dims = 0 : i64, device = "/device:CPU:0", validate_indices = true}: (tensor<*x!tf_type.resource>, tensor<?xi32>) -> (tensor<*xi32>)
  func.return %0 : tensor<*xi32>
}

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">
    (ArgumentList $params, $indices, $axis),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"batch_dims"> $batch_dims))),

def DecomposeResourceGather : Pat<
  (TF_ResourceGatherOp:$old_result
    $resource, $indices, $batch_dims, $validate_indices),
  (TF_GatherV2Op:$dest
    (CreateTFReadVariableOp $old_result, $old_result, $resource),
    $indices,
    (TF_ConstOp $batch_dims), // axis
    $batch_dims
  ), [], [(CopyXlaOutsideCompilationAttr $old_result, $dest)]>;

    if (batch_dims.size() != 1 || batch_dims[0] != 0) {
      return failure();
    }
    int64_t batch_dim = batch_dims[0];
    // Batch dim in operand and start indices should match.
    if (operand_type.getDimSize(batch_dim) > max_batch_size ||
        operand_type.getRank() != rank_two ||
        start_indices_type.getRank() != rank_two ||

// CHECK-NEXT:      %17 = "tfl.gather"(%14, %16) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32>
// CHECK-NEXT:      %18 = "tfl.pseudo_const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64>
// CHECK-NEXT:      %19 = "tfl.gather"(%15, %18) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<1xi64>) -> tensor<1xi32>