%52 = "tfl.concatenation"(%51, %0) {axis = 1 : i32, fused_activation_function = "NONE"} : (tensor<1x2x3xf32>, tensor<1x2x3xf32>) -> tensor<1x4x3xf32>
  %53 = "quantfork.stats"(%52) {layerStats = dense<[0.0, 1.0]> : tensor<2xf32>} : (tensor<1x4x3xf32>) -> tensor<1x4x3xf32>

  // CHECK-LABEL: func @set_bound
  func.func @set_bound(%arg0: tensor<i32>) -> tensor<i32> {
    %bound = "tf.Const"() {value = dense<16> : tensor<i32>} : () -> tensor<i32>

    // CHECK: %[[RESULT:.*]] = mhlo.custom_call @SetBound(%arg0) {backend_config = "", mhlo.literal = dense<16> : tensor<i32>}
    %bounded = "tf.XlaSetBound"(%arg0, %bound) : (tensor<i32>, tensor<i32>) -> tensor<i32>

    // CHECK: return %[[RESULT]]

//CHECK-DAG:    %[[c1:.+]] = stablehlo.constant dense<1> : tensor<i32>
//CHECK-DAG:    %[[c2:.+]] = stablehlo.constant dense<0> : tensor<i32>
//CHECK-NEXT:    %[[c3:.+]] = stablehlo.dynamic_update_slice %arg0, %[[c0]], %[[c1]], %[[c2]], %[[c2]] : (tensor<2x1x2xf32>, tensor<1x1x2xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2x1x2xf32>
//CHECK-NEXT:    return %[[c3:.+]] : tensor<2x1x2xf32>
//CHECK-NEXT:  }

}

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {
  %result = "tf.IfrtCall"(%arg0, %arg1) <{program_id = 1234 : i64, variable_arg_indices = [0 : i32, 1 : i32], variable_names = ["a", "b"]}> : (tensor<?xf32>, tensor<?xf32>) -> (tensor<1x1xf32>)
  func.return
}

// -----
func.func @test_ifrt_call_fail_unsorted_variable_arg_indices(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {

  %5 = "tfl.depthwise_conv_2d"(%2, %4, %cst) {depth_multiplier = 4 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<1x224x224x3xf32>, tensor<32x3x3x3xf32>, tensor<32xf32>) -> tensor<1x112x112x32xf32>

    %read0 = "tf.ReadVariableOp"(%arg0) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
    %read1 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
    %partitioned_input = "tf.TPUPartitionedInputV2"(%read0, %read1) {N = 2 : i64, partition_dims = []} : (tensor<i32>, tensor<i32>) -> tensor<i32>

func.func @main(%arg0: tensor<1x1024xf32>) -> tensor<1x3xf32> {
  %0 = "tf.Const"() <{value = dense<0.5> : tensor<1024x3xf32>}> : () -> tensor<1024x3xf32>
  %1:4 = "tf.CustomAggregator"(%arg0) <{id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {min = 7.547870e-07 : f32, max = 0.999992311 : f32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)

static constexpr char kPlatformName[] = "Host";
constexpr char kEntryFuncName[] = "main";

static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    func.return
  }
})";

MlirToHloArgs CreateTestMlirToHloArgs(const char* module_str = kMlirModuleStr) {
  MlirToHloArgs mlir_to_hlo_args;

  %2 = "tf.opC"() {is_stateless = true} : () -> tensor<f32>
  %3 = "tf.opD"(%ri_0, %ri_1, %arg2, %2) {_xla_compile_device_type = "TPU", _replication_info = "replicate", is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<i32>
  %ro_0:2 = "tf.TPUReplicatedOutput"(%3) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)

func.func @qi32_per_axis() -> tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>> {
  // CHECK-LABEL: @qi32_per_axis
  // CHECK: <{qtype = tensor<3x3x!quant.uniform<i32:f32:1, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>, value = dense<1> : tensor<3x3xi32>}> : () -> tensor<3x3x!quant.uniform<i32:f32:1, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>