%0, %1 = "tf.TopKV2"(%arg0, %arg1) : (tensor<8xf32>, tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>)
  func.return %0, %1: tensor<?xf32>, tensor<?xi32>

// CHECK-LABEL: topk
// CHECK:  "tfl.topk_v2"(%arg0, %arg1)
// CHECK:  return
}

func.func @topk_2(%arg0: tensor<8xf32>) -> (tensor<2xf32>, tensor<2xi32>) {
  %0 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>

  TRANSPOSE = 39,
  MEAN = 40,
  SUB = 41,
  DIV = 42,
  SQUEEZE = 43,
  UNIDIRECTIONAL_SEQUENCE_LSTM = 44,
  STRIDED_SLICE = 45,
  BIDIRECTIONAL_SEQUENCE_RNN = 46,
  EXP = 47,
  TOPK_V2 = 48,
  SPLIT = 49,
  LOG_SOFTMAX = 50,
  // DELEGATE is a special op type for the operations which are delegated to
  // other backends.
  // WARNING: Experimental interface, subject to change
  DELEGATE = 51,

// CHECK:           %[[VALUES:.*]], %[[INDICES:.*]] = "tfl.topk_v2"(%[[ARG]], %[[K]]) : (tensor<3x10xf32>, tensor<i32>) -> (tensor<3x5xf32>, tensor<3x5xi32>)
// CHECK:           return %[[VALUES]] : tensor<3x5xf32>
// CHECK:         }
func.func @optimizeTopKOnlyValues(%arg: tensor<3x10xf32>) -> tensor<3x5xf32>{
  %K = arith.constant dense<10> : tensor<i32>

  TRANSPOSE = 39,
  MEAN = 40,
  SUB = 41,
  DIV = 42,
  SQUEEZE = 43,
  UNIDIRECTIONAL_SEQUENCE_LSTM = 44,
  STRIDED_SLICE = 45,
  BIDIRECTIONAL_SEQUENCE_RNN = 46,
  EXP = 47,
  TOPK_V2 = 48,
  SPLIT = 49,
  LOG_SOFTMAX = 50,
  // DELEGATE is a special op type for the operations which are delegated to
  // other backends.
  // WARNING: Experimental interface, subject to change
  DELEGATE = 51,

  %0, %1 = "tfl.topk_v2"(%arg0, %arg1) : (tensor<8xf32>, tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>)
  func.return %0, %1: tensor<?xf32>, tensor<?xi32>
}

// -----

// CHECK-LABEL: topk
func.func @topk(%arg0: tensor<*xf32>, %arg1: tensor<i32>) -> (tensor<*xf32>, tensor<*xi32>) {
  %0, %1 = "tfl.topk_v2"(%arg0, %arg1) : (tensor<*xf32>, tensor<i32>) -> (tensor<*xf32>, tensor<*xi32>)

  TRANSPOSE = 39,
  MEAN = 40,
  SUB = 41,
  DIV = 42,
  SQUEEZE = 43,
  UNIDIRECTIONAL_SEQUENCE_LSTM = 44,
  STRIDED_SLICE = 45,
  BIDIRECTIONAL_SEQUENCE_RNN = 46,
  EXP = 47,
  TOPK_V2 = 48,
  SPLIT = 49,
  LOG_SOFTMAX = 50,
  // DELEGATE is a special op type for the operations which are delegated to
  // other backends.
  // WARNING: Experimental interface, subject to change
  DELEGATE = 51,

  func.return %0#0, %0#1: tensor<16x?xf32>, tensor<16x?xi32>
}

// -----

// CHECK-LABEL: topk_v2
// CHECK-SAME: %[[INPUT:.*]]: tensor<16x16xf32>
func.func @topk_v2(%input: tensor<16x16xf32>) -> (tensor<16x8xf32>, tensor<16x8xi32>) {
  %k = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>

// TODO(jpienaar): Check that input has one or more dimensions.
// TODO(jpienaar): Check that k is less or equal the internal dimension
def TFL_TopKV2Op: TFL_Op<"topk_v2", [
    QuantizableResult,
    Pure,
    TFL_OperandHasRankAtLeast<0, 1>,
    TFL_OperandHasRank<1, 0>,
    PredOpTrait<"result and input element type match",
      TFL_TCresVTEtIsSameAsOp<0,0>>,

*   `parallel_for`: Add converter for `MatrixDiag`.
*   Add `narrow_range` attribute to `QuantizeAndDequantizeV2` and V3.
*   Added new op: `tf.strings.unsorted_segment_join`.
*   Add HW acceleration support for `topK_v2`.
*   Add new `TypeSpec` classes.
*   CloudBigtable version updated to v0.10.0.
*   Expose `Head` as public API.
*   Update docstring for gather to properly describe the non-empty `batch_dims`
    case.

    "TRANSPOSE",
    "MEAN",
    "SUB",
    "DIV",
    "SQUEEZE",
    "UNIDIRECTIONAL_SEQUENCE_LSTM",
    "STRIDED_SLICE",
    "BIDIRECTIONAL_SEQUENCE_RNN",
    "EXP",
    "TOPK_V2",
    "SPLIT",
    "LOG_SOFTMAX",
    "DELEGATE",
    "BIDIRECTIONAL_SEQUENCE_LSTM",
    "CAST",
    "PRELU",
    "MAXIMUM",
    "ARG_MAX",
    "MINIMUM",
    "LESS",
    "NEG",