func.return
}

func.func private @write_chain_func(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {
  %cst = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
  %0 = "tf.XlaAllReduce"(%arg0, %cst) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<1xf32>, tensor<1x2xi32>) -> tensor<1xf32>
  return %0 : tensor<1xf32>

  %0 = "tfl.pseudo_sparse_qconst"() {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>, s_param = #tfl.sparsity_parameter<traversal_order = [0, 1, 2, 3], block_map...

  %0 = "tfl.pseudo_sparse_qconst"() {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>, s_param = #tfl.sparsity_parameter<traversal_order = [0, 1, 2, 3], block_map...

  %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>>

  %weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // CHECK: "tf.AddV2"
  // CHECK: mhlo.constant
  // CHECK-SAME{LITERAL}: dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>
  %0 = "tf.AddV2"(%input, %input) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>

  // CHECK: "mhlo.dot"
  // CHECK-SAME: (tensor<?x?xf32>, tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32>

  %group_assignment = "tf.Const"() { value = dense<[[0, 1]]> : tensor<1x2xi32> } : () -> tensor<1x2xi32>
  %instance_key = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
  %group_size, %group_key = "tf.CollectiveAssignGroupV2"(%group_assignment, %rank, %key_base) {} : (tensor<1x2xi32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
  // CHECK-NOT: "tf.CollectiveAssignGroupV2"

  %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>