// CHECK: %[[V0]] : tensor<17x17x24xf32>
}

// CHECK-LABEL: @broadcast_batch_matmul_v2_lhs
func.func @broadcast_batch_matmul_v2_lhs(%arg0: tensor<17x17xf32>, %arg1: tensor<17x17x24xf32>) -> tensor<17x17x24xf32> {
  %cst = arith.constant dense<[17, 17, 17]> : tensor<3xi64>
  %0 = "tf.BroadcastTo"(%arg0, %cst) : (tensor<17x17xf32>, tensor<3xi64>) -> tensor<17x17x17xf32>

    %recurrent_stats = "quantfork.stats"(%recurrent_input) {layerStats = dense<[-2.0, 1.0]> : tensor<2xf32>} : (tensor<1x20xf32>) -> tensor<1x20xf32>
    %cell_input = arith.constant dense<1.0> : tensor<1x20xf32>
    %cell_stats = "quantfork.stats"(%cell_input) {layerStats = dense<[-2.73090601, 7.94872093]> : tensor<2xf32>} : (tensor<1x20xf32>) -> tensor<1x20xf32>
    %0 = "tfl.unidirectional_sequence_lstm"(%arg0,

  %split_dim_2 = arith.constant dense<1> : tensor<1xi32>
  %4, %5 = "tfl.split"(%split_dim_2, %arg0) {num_splits = 2 : i32} : (tensor<1xi32>, tensor<16x4xf32>) -> (tensor<16x2xf32>, tensor<16x2xf32>)
  %6:2 = "tfl.split"(%split_dim_2, %arg0) {num_splits = 2 : i32} : (tensor<1xi32>, tensor<16x4xf32>) -> (tensor<16x2xf32>, tensor<16x?xf32>)

// quantize_composite_functions for module-level evaluation of functions.

module attributes {tf_saved_model.semantics} {
// CHECK: quantize_simple_xla_call_module(%[[ARG_0:.+]]: tensor<1x4xf32>)
  func.func private @quantize_simple_xla_call_module(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
    %0 = stablehlo.constant dense<1.000000e+00> : tensor<4x3xf32>

  %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x2xbf16>) -> tensor<1x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]])

  func.return %0 : tensor<1024xf32>

  // 1.0 * 2.0 * 512 + 4.0 = 1028.0
  // CHECK: %[[CST:.*]] = arith.constant dense<1.028000e+03> : tensor<1024xf32>
  // CHECK:  return %[[CST]]
}

// CHECK-LABEL: @ConstantFoldFullyConnectedNoBias
func.func @ConstantFoldFullyConnectedNoBias() -> tensor<1024xf32> {

// 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.

  %0 = "tf.Mean"(%arg0, %arg1) : (tensor<2x2xf32>, tensor<1xi32>) -> tensor<1x2xf32>
  func.return %0 : tensor<1x2xf32>

// CHECK-LABEL: mean
// CHECK:  "tfl.mean"(%arg0, %arg1) <{keep_dims = false}> : (tensor<2x2xf32>, tensor<1xi32>) -> tensor<1x2xf32>
}

func.func @mean_true(%arg0: tensor<2x2xf32>, %arg1: tensor<1xi32>) -> tensor<1x2xf32> {

  %10 = "tfl.concatenation"(%9, %7) {axis = -1 : i32, fused_activation_function = "NONE"} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<1x4xf32>
  %11 = "quantfork.stats"(%10) {layerStats = dense<[-0.488159984, 0.26483655]> : tensor<2xf32>} : (tensor<1x4xf32>) -> tensor<1x4xf32>
  %13 = "tfl.concatenation"(%9, %3) {axis = -1 : i32, fused_activation_function = "NONE"} : (tensor<1x2xf32>, tensor<1x1xf32>) -> tensor<1x3xf32>

// CHECK-DAG: %[[input_18:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-8.000000e-01, 1.600000e+00]> : tensor<2xf32>}> : (tensor<1x4xf32>) -> tensor<1x4xf32>
// 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]+}})