%dot = stablehlo.dot %q1, %q2 : (tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>, tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>) -> tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>
  %dq = "quantfork.dcast"(%dot) : (tensor<1x2x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x2xf32>
  func.return %dq : tensor<1x2xf32>
}

// -----

// CHECK-LABEL: @convert_quantfork_qdq_to_stablehlo_uniform_qdq

    %cst_1 = "tfl.pseudo_qconst"() {qtype = tensor<1x20x!quant.uniform<i8:f32, 0.006:-34>>, value = dense<1> : tensor<1x20xi8>} : () -> tensor<1x20x!quant.uniform<i8:f32, 0.006:-34>>
    %cst_2 = "tfl.no_value"() {value = unit} : () -> none

// CHECK-SAME: tensor<1x3xf32>, tensor<1x3xf32>,
// CHECK-SAME: none, none, none, none)
// CHECK-SAME: -> tensor<1x2x3xf32>
// CHECK: "tfl.quantize"(%[[lstm]]) <{qtype = tensor<1x2x3x!quant.uniform<i16:f32, {{.*}}>>}> {volatile} : (tensor<1x2x3xf32>) -> tensor<1x2x3x!quant.uniform<i16:f32, {{.*}}>>

}

// CHECK-LABEL: QuantizeUnidirectionalLstmFullPerAxis

    return %0 : tensor<1x3xf32>
  }
}

// CHECK-LABEL: quantize_dot_general_per_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>

  %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
  %1 = "tfl.pad"(%arg0, %0) : (tensor<15600xf32>, tensor<1x2xi32>) -> tensor<15600xf32>
  func.return %1 : tensor<15600xf32>
  // CHECK:  return %arg0
}

func.func @ConstFoldPadV2(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
  %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>

}

// -----

func.func @testConcat(%arg0: tensor<1x2xi32>, %arg1: tensor<1x2xi32>) -> tensor<2x2xi32> {
  // CHECK: "tfl.concatenation"(%arg0, %arg1) <{axis = 0 : i32, fused_activation_function = "NONE"}>
  %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2xi32>, tensor<1x2xi32>) -> tensor<2x2xi32>
  func.return %0 : tensor<2x2xi32>
}

// -----

        }) : () -> tensor<f32>
        "tf_device.cluster"() ({
          %group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
          %arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
          "tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
          tf_device.return
        }) : () -> ()

  // 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> {
    // CHECK: %[[Q1:.*]] = "quantfork.qcast"(%[[ARG0]]) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>

func.func @pack(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2x2xi32> {
  // CHECK: mhlo.reshape {{.*}} : (tensor<2xi32>) -> tensor<1x2xi32>
  // CHECK: mhlo.reshape {{.*}} : (tensor<2xi32>) -> tensor<1x2xi32>
  // CHECK: "mhlo.concatenate"({{.*}}) <{dimension = 0 : i64}> : (tensor<1x2xi32>, tensor<1x2xi32>) -> tensor<2x2xi32>

  %0 = "tf.Pack"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xi32>

func.func @testDontConvertPackToReshapeDynamicShape(%arg0: tensor<2x?xf32>) -> tensor<1x2x?xf32> {
  %0 = "tf.Pack"(%arg0) {axis = 0 : i64} : (tensor<2x?xf32>) -> tensor<1x2x?xf32>
  func.return %0 : tensor<1x2x?xf32>
  // CHECK: %[[PACK:.*]] = "tf.Pack"(%arg0) <{axis = 0 : i64}> : (tensor<2x?xf32>) -> tensor<1x2x?xf32>
  // CHECK: return %[[PACK]] : tensor<1x2x?xf32>
}

// CHECK-LABEL: while_with_id_passthrough