%3 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>
   %4 = "quantfork.dcast"(%3) : (tensor<1x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x2xf32>
// expected-error @+2 {{Failed to find a valid entry function}}

    %7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x2xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x2xi8>
    %8 = stablehlo.convert %7 : (tensor<1x2xi8>) -> tensor<1x2xf32>
    %9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
    %10 = stablehlo.dot_general %8, %9, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>

    return %1: tensor<1x3xbf16>
  }

  %3 = "tf.TPUExecute"(%2#0, %1, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
  %4 = "tf.TPUExecute"(%2#1, %1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>

  %0 = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>
  // CHECK: tfrt_fallback_async.const_dense_tensor  dense<1.000000e+00> : tensor<1xbf16>
  %1 = "tf.Const"() {device = "/device:CPU:0", value = dense<[1.0]> : tensor<1xbf16>} : () -> tensor<4xbf16>
  // CHECK: corert.executeop({{.*}}) "tf.Const"() {dtype = ui64, value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : 1

  func.return %6 : tensor<1x401408xf32>
}

// CHECK-LABEL: quantized_constant
func.func @quantized_constant(%arg0: tensor<1x2xf32>) -> tensor<2x2xf32> {
  %1 = "tfl.quantize"(%arg0) {qtype = tensor<1x2x!quant.uniform<u8:f32, 1.0>>, volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<u8:f32, 1.0>>

func.func @equal_broadcast(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> {
  // CHECK-DAG: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
  // CHECK-NEXT: mhlo.compare EQ, %[[LHS_BCAST]], %arg1
  %0 = "tf.Equal"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1>
  func.return %0: tensor<1x2xi1>
}

// CHECK-LABEL: func @equal_broadcast_no_incompatible_shapes_error

// CHECK-DAG: %[[input_19:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-2.000000e+00, 4.000000e+00]> : tensor<2xf32>}> : (tensor<1x2xf32>) -> tensor<1x2xf32>

// CHECK: "tfl.unidirectional_sequence_lstm"({{.*}}, %[[input_18]], %[[input_19]], %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}})

// RUN:   -split-input-file | FileCheck %s --check-prefix=STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY

// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: @main
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
  %0 = stablehlo.add %arg0, %arg0 : tensor<1x2xf32>
  return %0 : tensor<1x2xf32>
}
// Tests that `composite_add_fn_1` does not quantize when quantizing
// only compute-heavy ops.

// deferred.

// CHECK-LABEL: add_with_activation_transpose_rank_two
func.func @add_with_activation_transpose_rank_two(%arg0: tensor<1x2xf32>) -> tensor<2x1xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<2x1xf32>
  %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
  %2 = stablehlo.add %1, %0 : tensor<2x1xf32>
  return %2 : tensor<2x1xf32>
}