//CHECK-NEXT: return %0#0 : tensor<i64>
//CHECK-NEXT:}

func.func @gather(%operand: tensor<3x4x2xi32>, %start_indices: tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>{
  %result = "vhlo.gather_v1"(%operand, %start_indices) <{
    offset_dims = #vhlo.tensor_v1<dense<[2, 3]> : tensor<2xi64>>,
    collapsed_slice_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,

// 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]])

// -----

  };

  LogicalResult matchAndRewrite(
      mhlo::GatherOp gather_op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    if (succeeded(ConvertGatherOpToSlice(gather_op, rewriter))) {
      return success();
    }

    Value operand = gather_op.getOperand();
    Value start_indices = gather_op.getStartIndices();

    // Can only convert with static shaped gather.

  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 =

      gather_op.getDimensionNumbers().getOffsetDims().begin(),
      gather_op.getDimensionNumbers().getOffsetDims().end());
  std::vector<int64_t> collapsed_slice_dims_vec(
      gather_op.getDimensionNumbers().getCollapsedSliceDims().begin(),
      gather_op.getDimensionNumbers().getCollapsedSliceDims().end());
  std::vector<int64_t> start_index_map_vec(

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

  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  Operation* gather_op = FindOperationOfType<GatherOp>(main_fn);
  const auto spec = GetStableHloQuantConstraints(gather_op);

  EXPECT_THAT(spec, NotNull());
  EXPECT_THAT(spec->has_same_operand_and_result_type_requirement, IsTrue());
}

}  // namespace

        "//tensorflow/lite/tools/optimize:testdata/fc.bin",
        "//tensorflow/lite/tools/optimize:testdata/fc_qat.bin",
        "//tensorflow/lite/tools/optimize:testdata/gather_nd.bin",
        "//tensorflow/lite/tools/optimize:testdata/lstm_calibrated.bin",
        "//tensorflow/lite/tools/optimize:testdata/lstm_calibrated2.bin",
        "//tensorflow/lite/tools/optimize:testdata/lstm_quantized.bin",