// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor<f32>
}

// -----

module {
  func.func @quantize_xlagather(%arg0: tensor<10x2xi32>) -> tensor<1x300x10xf32> {
    %cst_0 = "tf.Const"() {device = "", value = dense<1.000000e+01> : tensor<200x100x300xf32>} : () -> tensor<200x100x300xf32>

    // CHECK:      metadata
    // CHECK-SAME: args
    // CHECK-SAME: shape
    // CHECK-SAME: dim
    // CHECK-SAME: size: 1
    // CHECK-SAME: dim
    // CHECK-SAME: size: 2
    // CHECK-SAME: dim
    // CHECK-SAME: size: 3

    func.return %0: tensor<1x2x3xi32>
  }

  }

  // CHECK-LABEL: batchmatmulv2
  func.func @batchmatmulv2(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
    // CHECK: mhlo.reduce
    // CHECK: mhlo.dot_general
    // CHECK: mhlo.transpose
    %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
    func.return %0 : tensor<3x4x4xf32>

  %0 = stablehlo.constant dense<1> : tensor<3x4x2xi32>
  %1 = stablehlo.constant dense<1> : tensor<2x3x2xi64>
  %2 = "stablehlo.gather"(%0, %1) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = array<i64: 1, 2, 2>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>

// CHECK-LABEL: squeezeToReshape
func.func @squeezeToReshape(%arg0: tensor<1x1x2xf32>) -> tensor<2xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<1x1x2xf32>) -> tensor<2xf32>
  func.return %0 : tensor<2xf32>

  // CHECK-DAG: [[CONST:.*]] = arith.constant dense<2> : tensor<1xi32>
  // CHECK: %[[RESULT:.*]] = "tfl.reshape"(%arg0, %[[CONST:.*]]) : (tensor<1x1x2xf32>, tensor<1xi32>) -> tensor<2xf32>
  // CHECK: return %[[RESULT]]
}

    return success();
  }
};

// Fuses Unpack with proceeding Concatenation to Reshape if output type has
// static shape and activation function is none. For example:
//
//   // %input: tensor<1x3x2xf32>
//   %unpack:3 = "tfl.unpack"(%input) {axis = 1 : i32, num = 3 : i32}
//   %res = "tfl.concatenation"(%unpack#0, %unpack#1, %unpack#2)
//        {axis = -1 : i32, fused_activation_function = "NONE"}
//

func.func @invalidSelect(%arg0: tensor<1x8xi1>, %arg1: tensor<1x8x8xi32>, %arg2: tensor<1x8x8xi32>) -> tensor<1x8x8xi32> {
  // expected-error @+1 {{requires that pred is a scalar OR has the same rank as t and e OR is a vector}}
  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<1x8xi1>, tensor<1x8x8xi32>, tensor<1x8x8xi32>) -> tensor<1x8x8xi32>
  func.return %0: tensor<1x8x8xi32>
}

// -----