// CHECK-NEXT:    return %2 : tensor<3x2xf32>
}

func.func @transpose_3d(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> {
  %0 = "mhlo.transpose"(%arg0) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32>
  func.return %0 : tensor<3x2x1xf32>

// CHECK-LABEL:   transpose_3d
// CHECK-NEXT:    %0 = "tfl.pseudo_const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64>

  constexpr absl::string_view kModuleGather = R"mlir(
    module {
      func.func @main() -> (tensor<2x3x2x2xf32>) {
        %0 = stablehlo.constant dense<1.0> : tensor<3x4x2xf32>
        %1 = stablehlo.constant dense<2> : tensor<2x3x2xi64>
        %2 = "stablehlo.gather"(%0, %1) {
          dimension_numbers = #stablehlo.gather<
            offset_dims = [2, 3],

func.func @no_qdq_except_conv_and_dot_general(%arg0: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> attributes {tf._original_func_name = "main_0"} {
  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<3x4x2xf32>} : () -> tensor<3x4x2xf32>
  %0 = "tf.XlaCallModule"(%cst, %arg0) {
    Sout = [#tf_type.shape<1x3>], _entry_function = @composite_gather_fn,
    _original_entry_function = "composite_gather_fn", _quantization_method = "weight_only_ptq { }",

}

// -----

// CHECK-LABEL: topk_2
func.func @topk_2(%arg0: tensor<3x4x8xf32>) -> (tensor<3x4x2xf32>, tensor<3x4x2xi32>) {
  %0 = arith.constant dense<2> : tensor<i32>
  %1:2 = "tfl.topk_v2"(%arg0, %0) : (tensor<3x4x8xf32>, tensor<i32>) -> (tensor<3x4x2xf32>, tensor<3x4x2xi32>)
  func.return %1#0, %1#1: tensor<3x4x2xf32>, tensor<3x4x2xi32>
}

// -----

// CHECK-LABEL: topk_d

  %gather = "tf.TensorListGather"(%tl, %arg1, %elem_shape) : (tensor<!tf_type.variant<tensor<8x9xf32>>>, tensor<3xi32>, tensor<2xi32>) -> tensor<3x8x9xf32>
  // CHECK-NEXT: return %[[GATHER]] : tensor<3x8x9xf32>
  func.return %gather: tensor<3x8x9xf32>
}

// -----

// Test scatter into existing tensor list.

  %0:3 = "tf.Unpack"(%input) {axis = 1} : (tensor<4x3x6xf32>) -> (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>)
  // return %[[RES1]], %[[RES2]], %[[RES3]]
  func.return %0#0, %0#1, %0#2 : tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>
}

// -----

// CHECK-LABEL: func @unpack_dynamic
func.func @unpack_dynamic(%arg0: tensor<?x?x2xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {

  %0:3 = "tf.Split"(%arg0, %arg1) : (tensor<i32>, tensor<1x4x3x3xf32>) -> (tensor<1x4x3xf32>, tensor<1x4x3xf32>, tensor<1x4x3xf32>)
  func.return %0#0 : tensor<1x4x3xf32>

  // CHECK-LABEL: split
  // CHECK: "tfl.split"(%arg0, %arg1) <{num_splits = 3 : i32}> : (tensor<i32>, tensor<1x4x3x3xf32>) -> (tensor<1x4x3xf32>, tensor<1x4x3xf32>, tensor<1x4x3xf32>)
}

func.func @foldFill() -> (tensor<3x2x1xf32>, tensor<*xf32>, tensor<*xcomplex<f32>>) {
  %0 = "tf.Const"() {value = dense<[3, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32>
  %1 = "tf.Const"() {value = dense<23.0> : tensor<f32>} : () -> tensor<f32>
  // CHECK-DAG: "tf.Const"() <{value = dense<2.300000e+01> : tensor<3x2x1xf32>}>
  %2 = "tf.Fill"(%0, %1) : (tensor<3xi32>, tensor<f32>) -> tensor<3x2x1xf32>

}

// -----

// CHECK-LABEL: tensorlistConcat
func.func @tensorlistConcat(%arg0: tensor<3x2x2xf32>, %lead: tensor<i64>) -> (tensor<?x2xf32>, tensor<0xi64>) {
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %list = "tf.TensorListFromTensor"(%arg0, %cst) : (tensor<3x2x2xf32>, tensor<2xi32>) -> tensor<!tf_type.variant<tensor<f32>>>

//                (tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>)
//
// will be converted into:
//
//   %0 = "mhlo.slice"(%input) {
//             limit_indices = dense<[4, 2]> : tensor<2xi64>,
//             start_indices = dense<0> : tensor<2xi64>,
//             strides = dense<1> : tensor<2xi64>} :
//        (tensor<4x6xf32>) -> tensor<4x2xf32>