%14 = mhlo.compare  LT, %13, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
  %15 = mhlo.add %13, %4 : tensor<32x32xi64>
  %16 = mhlo.select %14, %15, %13 : tensor<32x32xi1>, tensor<32x32xi64>
  %17 = mhlo.reshape %16 : (tensor<32x32xi64>) -> tensor<32x32x1xi64>
  %18 = "mhlo.concatenate"(%12, %17) <{dimension = 2 : i64}> : (tensor<32x32x1xi64>, tensor<32x32x1xi64>) -> tensor<32x32x2xi64>

}

func.func @matrix_set_diag_v3_non_zero_k(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi32>) -> tensor<3x3xi64> {
  %cst = arith.constant dense<1> : tensor<i32>
  %0 = "tf.MatrixSetDiagV3"(%arg0, %arg1, %cst) : (tensor<3x3xi64>, tensor<3xi32>, tensor<i32>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

// CHECK-LABEL: @matrix_set_diag_v3_non_zero_k
// CHECK: tf.MatrixSetDiagV3
}

    // CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[CALL]]

    // CHECK: %[[REDUCE:.*]] = "stablehlo.reduce_window"(%[[RESHAPE]], %[[Q0]])
    // CHECK{LITERAL}: padding = dense<[[0, 0], [1, 1], [0, 0]]> : tensor<3x2xi64>
    // CHECK-SAME: window_dimensions = array<i64: 1, 3, 1>
    // CHECK: %[[ARG1:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>, %[[ARG2:.*]]: tensor<!quant.uniform<i8:f32, 3.000000e-01:1>>

    tensor<3x4x2x2x!quant.uniform<i8:f32, 3.000000e-01:-5>>,
    tensor<2x3x2xi64>
  ) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 3.000000e-01:-5>>
  return %0 : tensor<2x3x2x2x!quant.uniform<i8:f32, 3.000000e-01:-5>>
}

// CHECK-LABEL: func @gather
// CHECK-SAME: %[[ARG_0:.+]]: tensor<3x4x2x2x!quant.uniform<i8:f32, 3.000000e-01:-5>>, %[[ARG_1:.+]]: tensor<2x3x2xi64>
// CHECK: %[[GATHER:.+]] = "tfl.gather_nd"(%[[ARG_0]], %[[ARG_1]])

  // CHECK-SAME: %[[ARG0:.*]]: tensor<3x4x5xf32>
  // CHECK-SAME: %[[ARG1:.*]]: tensor<3x5x2xf32>
  // CHECK-SAME: %[[ARG2:.*]]: tensor<2x3x2xi64>
  func.func private @composite_and_gather(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>, %arg2: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> {
    // CHECK: %[[Q1:.*]] = "quantfork.qcast"(%[[ARG0]]) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>

// -----

// Test that q/dq pair is not inserted between constant and XlaCallModule op
// whose entry function name does not include conv nor dot_general.

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) {

// CHECK-SAME:                      %[[VAL_1:.*]]: tensor<f32>) -> tensor<11x131xf32> {
// CHECK:           %[[VAL_2:.*]] = arith.constant dense<{{\[\[}}1, 2], [0, 3]]> : tensor<2x2xi64>
// CHECK:           %[[VAL_3:.*]] = "tf.PadV2"(%[[VAL_0]], %[[VAL_2]], %[[VAL_1]]) : (tensor<8x128xf32>, tensor<2x2xi64>, tensor<f32>) -> tensor<11x131xf32>
// CHECK:           return %[[VAL_3]] : tensor<11x131xf32>
// CHECK:         }

}

// CHECK-LABEL: @testMatrixSetDiag
func.func @testMatrixSetDiag(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi64>) -> tensor<3x3xi64> {
  %0 = "tf.MatrixSetDiag"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x3xi64>, tensor<3xi64>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

  // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}>

  %3 = "tfl.pseudo_qconst"() {qtype = tensor<32x12x!quant.uniform<u4<1:15>:f32, 0.395998308:8>>, value = dense<-7> : tensor<32x12xi4>} : () -> tensor<32x12x!quant.uniform<u4<1:15>:f32, 0.395998308:8>>
  %4 = "tfl.dequantize"(%3) : (tensor<32x12x!quant.uniform<u4<1:15>:f32, 0.395998308:8>>) -> tensor<32x12xf32>

  %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>
  func.return %2: tensor<2x3x2x2xi32>