}> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    func.return %0: tensor<1x2x2x3xf32>
  }

  func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> {
    return %arg0 : tensor<1x2x2x3xf32>
  }
}

// CHECK: module
// CHECK-NOT: tf.PartitionedCall
// CHECK-NOT: some_func

    has_token_input_output = false, module = "", platforms = [],
    version = 5 : i64
  } : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
  return %0 : tensor<1x2x2x2xf32>
}

// CHECK-LABEL: func.func @qdq_for_conv_weight_empty
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>

  %5 = "mhlo.broadcast_in_dim"(%4) {broadcast_dimensions = dense<[2, 3]> : tensor<2xi64>} : (tensor<6x6xf32>) -> tensor<1x3x6x6xf32>
  %6 = mhlo.divide %3, %5 : tensor<1x3x6x6xf32>
  return %6 : tensor<1x3x6x6xf32>
}

// CHECK-LABEL:   func.func @avg_pool2d_2(
// CHECK-SAME:                               %[[VAL_0:.*]]: tensor<1x3x6x6xf32>) -> tensor<*xf32> {
// CHECK:           %[[VAL_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>

// CHECK: stablehlo.maximum
// CHECK: (tensor<1x5x5x4xf32>, tensor<f32>) -> tensor<1x2x2x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[REDUCE_WINDOW_MAX]], dims = [0, 3, 1, 2] : (tensor<1x2x2x4xf32>) -> tensor<1x4x2x2xf32>
// CHECK: return %[[TRANSPOSE_1]]

// -----

// Tests that a `maximum(add(convolution(%activation, %weight), %bias), %zero)`

    rhs_quantization_max_val = 127 : i64
  } : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x2x1x2xf32>
  func.return %0 : tensor<1x2x1x2xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantized_convolution_hybrid_explicit
func.func @uniform_quantized_convolution_hybrid_explicit(%input: tensor<1x2x2x3xf32>) -> tensor<1x3x3x2xf32> {

  %2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  return %2 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @composite_conv2d_with_bias_and_relu6_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"

// -----

// Check the IfOp is set to stateful.

  %3 = "tfl.quantize"(%2) {qtype = tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>, volatile} : (tensor<1x2x1x3xf32>) -> tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>
  %5 = "tfl.dequantize"(%arg0) : (tensor<1x2x1x3x!quant.uniform<i8:f64, 1.0>>) -> tensor<1x2x1x3xf32>
  "tfl.assign_variable"(%1, %5) : (tensor<!tf_type.resource>, tensor<1x2x1x3xf32>) -> ()

    : (tensor<?x3x2x1xf32>, tensor<2x1x1x1xf32>) -> tensor<?x2x2x1xf32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xf32>, tensor<?x2x2x1xf32>) -> tensor<?x2x2x1xf32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x2x2x1xf32>, tensor<1xf32>) ->tensor<?x2x2x1xf32>
  return %2 : tensor<?x2x2x1xf32>
}

// -----

// CHECK-LABEL: func @dot_general_add_add

func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

  %3 = "tfl.quantize"(%2) {qtype = tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>, volatile} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>
  %4 = "tfl.dequantize"(%3) : (tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>) -> tensor<1x2x2x3xf32>