auto y = GatherV2(scope_, x, {9, 8, 7, 6}, /*axis=*/3);
  TensorShape y_shape({2, 3, 7, 4});
  RunTest(x, shape, y, y_shape);
}

TEST_F(ArrayGradTest, GatherV2Grad_BatchDim) {
  TensorShape shape({2, 100, 3});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
  GatherV2::Attrs attrs;
  attrs.batch_dims_ = 1;
  auto y =

// CHECK-NEXT: [[res:%.*]] = tfrt_fallback_async.executeop key(1) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.GatherV2"([[value]], {{.*}}, [[const]])
// CHECK-NEXT: tfrt.return [[out_chain]], [[res]] : !tfrt.chain, !tfrt_fallback.tf_tensor
func.func @gather(%indices: tensor<?xi32>,

 ExpandDims 1
 ReverseSequence 1
 Shape 1
 StridedSlice 1
 Transpose 3
cluster 6 size 11
 Cast 1
 Const 5
 GatherV2 1
 Shape 1
 StridedSlice 1
 Transpose 1
 ZerosLike 1
cluster 7 size 33
 All 2
 Cast 1
 Const 17
 Equal 2
 ExpandDims 2
 GatherV2 1
 Identity 1
 ReverseSequence 1
 Shape 2
 StridedSlice 2
 Transpose 2
cluster 8 size 11
 Cast 1
 Const 4

// CHECK-SAME: value = dense<1> : tensor<1xi32>
// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"()
// CHECK-SAME: value = dense<1> : tensor<1xi32>

// Check that the constants used for tf.GatherV2 and tf.AddV2 are different.
// CHECK: %[[GATHER_V2:.*]] = "tf.GatherV2"(%[[ARG_0]], %[[ARG_1]], %[[CST_1]])
// CHECK: %[[ADDV2:.*]] = "tf.AddV2"(%[[CST_0]], %[[CST_0]])

// -----

// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.Const
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.AddV2
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf.GatherV2
// CHECK-NEXT: tf.Mul
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf_mlrt.tf_await
// CHECK-NEXT: tf.TensorListSetItem

  }
}


// -----
//  Decomposable Resource Ops usage
//
// CHECK-LABEL:  func.func @serving_default
// CHECK:      "tf.VarHandleOp"
// CHECK-NEXT: "tf.IfrtLoadVariable"
// CHECK-NEXT: "tf.GatherV2"
// CHECK-NEXT:  return 
//
module {
  func.func @serving_default() -> tensor<1x3xbf16> {
    %cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>

                         %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>

  %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>

  if (mlir::isa<mlir::TF::StringType>(key_type.getElementType()))
    cost *= kLookupTableFindStringKeyCostScale;

  return cost;
}

// The cost function for tf.GatherV2.
int64_t InferGatherV2Cost(const CostContext& context, mlir::TF::GatherV2Op op) {
  return InferTensorSize(
      context, mlir::cast<mlir::TensorType>(op.getOutput().getType()));
}

  std::set<string> constrained_kernel = {"Diag",
                                         "DiagPart",
                                         "GatherNd",
                                         "GatherV2",
                                         "QuantizeAndDequantizeV2",
                                         "QuantizeAndDequantizeV3",
                                         "QuantizeAndDequantizeV4",