}
  func.func private @composite_gather_fn_1(%arg0: tensor<1024x3x4x3xf32>, %arg1: tensor<1xi32>, %arg2: tensor<i32>) -> tensor<1x3x4x3xf32> attributes {tf_quant.composite_function} {
    %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>
  }

  %cst = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %cst_0 = "tf.Const"() {device = "", value = dense<1.0> : tensor<128x32xf32>} : () -> tensor<128x32xf32>
  %0 = "tf.GatherV2"(%cst_0, %arg0, %cst) {batch_dims = 0 : i64, device = ""} : (tensor<128x32xf32>, tensor<6xi64>, tensor<i32>) -> tensor<6x32xf32>
  return %0 : tensor<6x32xf32>

// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<i32>

                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      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>

    // expected-remark@+1 {{Cost: 1}}
    %0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>
    // expected-remark@+1 {{Cost: 300}}
    %1 = "tf.GatherV2"(%arg0, %arg1, %0) : (tensor<1x2x20xf32>, tensor<3x5xi32>, tensor<1xi32>) -> tensor<1x3x5x20xf32>
    // expected-remark@+1 {{Cost: 40}}
    func.return %1 : tensor<1x3x5x20xf32>
}

  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.");
  }

  // Handle axis by transposing the axis dimension to be the first non-batch
  // dimension, compute the gradient and transpose the result back.

                         %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8>
      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>

      (TF_AssignVariableOp $v_resource, $new_v, (CreateConstBoolAttrFalse))
    ]
  >;

// Pattern to decompose tf.ResourceGather into tf.ReadVariable and tf.GatherV2.
def DecomposeResourceGather : Pat<
  (TF_ResourceGatherOp:$old_result
    $resource, $indices, $batch_dims, $validate_indices),
  (TF_GatherV2Op:$dest
    (CreateTFReadVariableOp $old_result, $old_result, $resource),

    %cst_1 = "tf.Const"() {value = dense<1.1> : tensor<3x3x1024x1xf32>} : () -> tensor<3x3x1024x1xf32>
    %cst_2 = "tf.Const"() {value = dense<1.1> : tensor<1024x3x4x3xf32>} : () -> tensor<1024x3x4x3xf32>
    %0 = "tf.GatherV2"(%cst_2, %arg0, %cst) {batch_dims = 0 : i64, device = ""} : (tensor<1024x3x4x3xf32>, tensor<1xi32>, tensor<i32>) -> tensor<1x3x4x3xf32>

    %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<*xi32>>>

    // CHECK-DAG: [[READVAR:%.+]] = "tf.ReadVariableOp"([[VAR]])
    // CHECK: [[GATHER:%.+]] = "tf.GatherV2"([[READVAR]], [[INDEX]], [[ZERO]]) <{batch_dims = 0 : i64}> {_xla_outside_compilation = "0"} : (tensor<*xi32>, tensor<?xi32>, tensor<i64>) -> tensor<*xi32>
    // CHECK: return [[GATHER]]

          (MHLO_FftOp $input, MHLO_FftTypeValue<"IFFT">, (GetInnerDimFromValue $res)),
          [(CheckInnerDimStatic $input)]>;

//===----------------------------------------------------------------------===//
// GatherV2 op patterns.
//===----------------------------------------------------------------------===//

// Here, $params and $indices needs to be ranked so that $axis and $batch_dims