// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
  // CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
  // CHECK-SAME: %[[ARG2:.*]]: tensor<1x3xi1>
  // CHECK-SAME: %[[ARG3:.*]]: tensor<1x3xf32>
  func.func private @composite_and_select(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<1x3xi1>, %arg3: tensor<1x3xf32>) -> tensor<1x3xf32> {

}

// -----

module {
  // CHECK-LABEL: dot_general
  func.func private @dot_general(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
    // CHECK: %[[q_weight:.*]] = "quantfork.qcast"
    // CHECK-SAME: -> tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {0.049663885371891529,0.060200210631363035}>>
    // CHECK: %[[dq_weight:.*]] = "quantfork.dcast"(%[[q_weight]])

  func.func @xla_spmd_full_to_shard_shape(%arg0: tensor<2x2xi64>) -> (tensor<1x2xi64>) {
    // CHECK: %[[SHARDING:.*]] = mhlo.custom_call @Sharding(%arg0) {backend_config = "", mhlo.sharding = "{devices=[2,1]0,1}"} : (tensor<2x2xi64>) -> tensor<2x2xi64>
    // CHECK: %[[MANUAL:.*]] = mhlo.custom_call @SPMDFullToShardShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{manual}"} : (tensor<2x2xi64>) -> tensor<1x2xi64>
    // CHECK: return %[[MANUAL]]

// CHECK: return %[[QUANT]]

// -----

// Tests that the pattern doesn't match when the input tensor's type is a
// quantized type.

func.func @uniform_quantize_op_quantized_input(%arg: tensor<2x2x!quant.uniform<i8:f32, 2.000000e+0:16>>) -> tensor<2x2x!quant.uniform<i8:f32, 3.000000e+0:127>> {

// TF-DAG: %[[cc_0:.*]] = "tf.ConcatV2"(%[[pc_1]], %[[pc_5]], %[[cst_0]]) : (tensor<2x2xi8>, tensor<2x2xi8>, tensor<i32>) -> tensor<2x4xi8>
// TF-DAG: %[[pc_7:.*]] = "tf.PartitionedCall"(%[[cc_0]], %[[cst_6]], %[[cst_5]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}
// TF-DAG: return %[[pc_7]]

// CHECK-LABEL: @all_keep_dim
func.func @all_keep_dim(%input: tensor<4x8xi1>) -> tensor<4x1xi1> {
  // CHECK: mhlo.reshape %{{.*}} : (tensor<4xi1>) -> tensor<4x1xi1>
  %dims = "tf.Const"() { value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
  %0 = "tf.All"(%input, %dims) {keep_dims = true} : (tensor<4x8xi1>, tensor<1xi32>) -> tensor<4x1xi1>
  func.return %0 : tensor<4x1xi1>
}

// -----

// CHECK-LABEL: @all_dynamic

                                             precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
                                             rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
                                             window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,

  %1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32> loc("XlaConv/feature_group_count")
  %2 = "tf.Const"() {value = dense<1> : tensor<2x2xi32>} : () -> tensor<2x2xi32> loc("XlaConv/padding")
  %3 = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32> loc("XlaConv/window_strides")

    %1 = "tf.Const"() {value = dense<2> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
    %2 = "tf.Identity"(%1) {device = ""} : (tensor<2x2xi32>) -> tensor<2x2xi32>
    // CHECK: SpaceToBatchND{{.*}}-> tensor<4x98x130x128xf32>
    %3 = "tf.SpaceToBatchND"(%arg0, %0, %2) {device = ""} : (tensor<1x192x256x128xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x?x?x128xf32>
    func.return
  }

//
//   %indices0 = "tf.Const"() {value = dense<4> : tensor<i32>}
//   %indices1 = "tf.Const"() {value = dense<[[3, 2], [1, 0]]> :
//   tensor<2x2xi32>} %0 = "tf.DynamicStitch"(%indices0, %indices1, %arg0,
//   %arg1)
//     : (tensor<i32>, tensor<2x2xi32>, tensor<2xf32>, tensor<2x2x2xf32>)
//     -> tensor<5x2xf32>
//
// is lowered to
//
//   %shape = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi64>}