// Run optimize-batch-matmul pass only and check the results.
// RUN: tf-opt %s -tfl-optimize-batch-matmul | FileCheck %s

// CHECK-LABEL: FuseTransposeFCRhsToBatchMatmul
func.func @FuseTransposeFCRhsToBatchMatmul(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xf32>, %arg2: none) -> tensor<16x128xf32> {
  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.transpose"(%arg1, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>

    precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>}> : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
  return %0 : tensor<1x1x64xf32>
}

//CHECK:func.func private @dot_general(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) -> tensor<1x1x64xf32> {
//CHECK-NEXT: %0 = "vhlo.dot_general_v1"(%arg0, %arg1) <{
//CHECK-SAME:    lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,

%4, %4) {_tflite_input_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 18, 19], device = ""} : (tensor<28x1x28xf32>, tensor<16x28xf32>, tensor<16x28xf32>, tensor<16x28xf32>, tensor<16x28xf32>, tensor<16x16xf32>, tensor<16x16xf32>, tensor<16x16xf32>, tensor<16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1x16xf32>, tensor<1x16xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<28x1x16xf32>)...

// CHECK:           }) : (tensor<16x1504xf32>, tensor<1xi32>, tensor<16xf32>) -> tensor<16x1504xf32>
// CHECK:           return %[[VAL_3]] : tensor<16x1504xf32>
// CHECK:         }
func.func @convert_scatter_update_reshape_indices_and_updates(
  %arg0: tensor<16x1504xf32>,
  %arg1: tensor<1xi32>,
  %arg2: tensor<16xf32>) -> tensor<16x1504xf32>
{
  %0 = "mhlo.scatter"(%arg0, %arg1, %arg2) ({

    %elem_shape = "tf.Const"() {value = dense<[-1, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
    %tensors = "tf._SomeOp"() : () -> tensor<10x16x1xf32>
    // CHECK: TensorListFromTensor
    // CHECK-SAME: tensor<!tf_type.variant<tensor<16x1xf32>>>
    %tl = "tf.TensorListFromTensor"(%tensors, %elem_shape) : (tensor<10x16x1xf32>, tensor<2xi32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>
    func.return
  }

// -----

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

  // CHECK:     chlo.top_k(%[[INPUT]], k = 8)
  %0:2 = "tf.TopKV2"(%input, %k): (tensor<16x16xf32>, tensor<i32>) -> (tensor<16x8xf32>, tensor<16x8xi32>)