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

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

// CHECK:  "tfl.gather"(%arg0, %arg1) <{axis = -1 : i32, batch_dims = 0 : i32}> : (tensor<1x2x20xf32>, tensor<3x5xi32>) -> tensor<1x2x3x5xf32>
}

func.func @gatherWithBatchDims(%arg0 : tensor<2x3x6xf32>, %arg1 : tensor<2x5xi32>) -> tensor<2x5x3x6xf32> {
  %0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>

  }

  int64_t batch_dims_i = batch_dims.getInt();
  if (batch_dims_i < 0) {
    batch_dims_i += indices_rank;
  }

  if (batch_dims_i > axis_i) {
    emitError(result.location,
              "axis should be bigger than or equal to batch_dims");
  }
  if (batch_dims_i >= params_rank || batch_dims_i > indices_rank) {
    emitError(result.location,

  %0 = "tf.Const"() { value = dense<[-1]> : tensor<1xi32> } : () -> tensor<1xi32>
  %1 = "tf.GatherV2"(%arg0, %arg1, %0) {batch_dims = -1 : i64} : (tensor<16x2x3xf32>, tensor<16x5xi32>, tensor<1xi32>) -> tensor<16x2x5x3xf32>
  func.return %1 : tensor<16x2x5x3xf32>
}

// -----

// Verify that the batch_dims can be equal to the rank of the indices.
func.func @testGatherV2(%arg0: tensor<16x4xf32>, %arg1: tensor<16xi32>) -> tensor<16xf32> {

}

// CHECK-LABEL: testGatherToV2
// Ensures that axis param and batch_dims attr use their default values of 0.
func.func @testGatherToV2(%params: tensor<4x3xf32>, %indices: tensor<1x2xi32>) -> tensor<2x3xf32> {
  // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

  %30 = "tf.GatherV2"(%9, %29, %5) {batch_dims = 0 : i64, device = ""} : (tensor<32x1xf32>, tensor<2x16xi32>, tensor<i32>) -> tensor<2x16x1xf32>
  %31 = "tf.AddV2"(%19, %28) {device = ""} : (tensor<2x16xi32>, tensor<2x16xi32>) -> tensor<2x16xi32>
  %32 = "tf.GatherV2"(%9, %31, %5) {batch_dims = 0 : i64, device = ""} : (tensor<32x1xf32>, tensor<2x16xi32>, tensor<i32>) -> tensor<2x16x1xf32>

    TFL_TensorOf<[I16, I32, I64]>:$indices,
    I32Attr:$axis,
    DefaultValuedOptionalAttr<I32Attr, "0">:$batch_dims
  );

  let builders =
  [
    OpBuilder<(ins "Value":$params, "Value":$indices, "IntegerAttr":$axis, "IntegerAttr":$batch_dims),
    [{ BuildGatherOp(&$_builder, $_state, params, indices, axis, batch_dims); }]>
  ];

  let results = (outs

    %indices = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
    %axis = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<i32>
    %1 = "tf.GatherV2"(%0, %indices, %axis) {batch_dims = 0 : i64, device = ""} : (tensor<3xi32>, tensor<2xi32>, tensor<i32>) -> tensor<2xi32>

    %last_dim = "tf.Const"() {value = dense<7> : tensor<1xi32>} : () -> tensor<1xi32>

}

// -----

// CHECK-LABEL: testGatherWithBatchDims
func.func @testGatherWithBatchDims(%arg0 : tensor<2xf32>, %arg1 : tensor<2xi32>) -> tensor<2xf32> {
  %0 = "tfl.gather"(%arg0, %arg1) {axis = 1 : i32, batch_dims = 2 : i32}: (tensor<2xf32>,tensor<2xi32>) -> tensor<2xf32>
  func.return %0 : tensor<2xf32>
}

// -----

// CHECK-LABEL: testGatherNd