func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.
  %0 = stablehlo.constant dense<0.0> : tensor<12xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<12xf32>) -> tensor<3x4xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<3x4xf32>) -> tuple<tensor<3x4xf32>>
  func.return %2 : tuple<tensor<3x4xf32>>
}


// CHECK-LABEL:   func @tfl_wrapped_jax_random_uniform(

module attributes {tf.versions = {producer = 888 : i32}} {
func.func @testReshapeShapeInference(%arg0: tensor<3x4xi32>) -> tensor<*xi32> {
  %cst = arith.constant dense<[1, 6, 2]> : tensor<3xi32>
  // CHECK: "tfl.reshape"(%arg0, %cst) : (tensor<3x4xi32>, tensor<3xi32>) -> tensor<1x6x2xi32>
  %0 = "tfl.reshape"(%arg0, %cst) : (tensor<3x4xi32>, tensor<3xi32>) -> tensor<*xi32>
  func.return %0 : tensor<*xi32>
}
}

// -----

  // CHECK-SAME: -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  // CHECK: %[[RES_INT:.*]] = mhlo.bitcast_convert %[[RES]] : (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>) -> tensor<3x2xi32>
  // CHECK: return %[[RES_INT]] : tensor<3x2xi32>

  %0 = "tf.UniformQuantizedAdd"(
    %arg0, %bias,
    %input_scales, %input_zps,
    %bias_scales, %bias_zps,
    %output_scales, %output_zps) {

  // attribute `_quantization_method` set to `"no_quantization {}"`.
  constexpr absl::string_view kXlaCallModuleOpWithQuantizationMethodAttr =
      R"mlir(
    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {

  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>

  // CHECK-NEXT:     %[[GROUP:.*]] = "tf.Const"()
  // CHECK-SAME:       [0, 1, 2, 3]
  // CHECK-NEXT:     %[[REDUCED:.*]] = "tf.XlaAllReduce"(%[[ID]], %[[GROUP]]) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x4xi32>) -> tensor<f32>
  // CHECK-NEXT:     "tf.OpA"(%[[REDUCED]]) : (tensor<f32>) -> ()
  tf_device.replicate {n = 4 : i32} {
    "tf_device.cluster"() ({
      "tf.OpA"(%arg0) : (tensor<f32>) -> ()
      tf_device.return

}

// CHECK-LABEL: testRedundantReshape
func.func @testRedundantReshape(%arg0: tensor<4x4xi32>) -> tensor<2x8xi32> {
  %0 = "tf.Const"() {value = dense<[8, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
  %1 = "tf.Const"() {value = dense<[2, 8]> : tensor<2xi32>} : () -> tensor<2xi32>
  %2 = "tf.Reshape"(%arg0, %0) : (tensor<4x4xi32>, tensor<2xi32>) -> tensor<8x2xi32>

// CHECK-LABEL: QuantizePad
func.func @QuantizePad(tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>, tensor<3x2xi32>) -> tensor<?xf32> {
^bb0(%arg0: tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>, %arg1: tensor<3x2xi32>):
  %0 = "tfl.dequantize"(%arg0) : (tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>) -> tensor<2x1x3xf32>
  %1 = "tfl.pad"(%0, %arg1) : (tensor<2x1x3xf32>, tensor<3x2xi32>) -> tensor<?xf32>
  func.return %1 : tensor<?xf32>

  func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<1x3xf32>, %arg1: tensor<3x64xf32>, %arg2: tensor<1x64xf32>) -> tensor<1x64xf32> {
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x64xf32>
    %1 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf32>, tensor<3x64xf32>) -> tensor<1x64xf32>
    %2 = stablehlo.add %1, %arg2 : tensor<1x64xf32>

  func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi32> {
    %0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi32>
    return %0 : tensor<1x2xi32>
  }
// CHECK: @uniform_quantize_0
  func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {