// CHECK-SAME: %[[ARG1:.*]]: tensor<*x!tf_type.resource<tensor<f32>>> {tf.device = "/device:TPU:1"},
  // CHECK-SAME: %[[ARG2:.*]]: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>> {tf.device = "/device:TPU:0"},
  // CHECK-SAME: %[[ARG3:.*]]: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>> {tf.device = "/device:TPU:1"})
  func.func @main(%arg0: !tf_res_f32 {tf.device = "/device:TPU:0"},
             %arg1: !tf_res_f32 {tf.device = "/device:TPU:1"},

}

// -----

// CHECK-LABEL: testSliceConcat
func.func @testSliceConcat(%arg0: tensor<3x1x512xf32>) -> tensor<3x1x512xf32> {
  %0 = "mhlo.slice"(%arg0) <{limit_indices = dense<[1, 1, 512]> : tensor<3xi64>, start_indices = dense<[0, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}> : (tensor<3x1x512xf32>) -> tensor<1x1x512xf32>

  %8 = "quantfork.stats"(%7) {
    layerStats = dense<[0.0, 1.4]> : tensor<2xf32>
  } : (tensor<?x1x1x3xf32>) -> tensor<?x1x1x3xf32>
  %9 = tfl.mul %8, %4 {fused_activation_function = "NONE"} : tensor<?x1x1x3xf32>
  %10 = "quantfork.stats"(%9) {
    layerStats = dense<[0.000000e+00, 5.6]> : tensor<2xf32>
  } : (tensor<?x1x1x3xf32>) -> tensor<?x1x1x3xf32>
  func.return %10 : tensor<?x1x1x3xf32>

  %cst_0 = "tf.Const"() {value = dense<0.400000e+00> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
  %cst_1 = "tf.Const"() {value = dense<0.200000e+00> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
  %cst_2 = "tf.Const"() {value = dense<0.300000e+00> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>

// MinElement-LABEL: QuantizeCustomOp
func.func @QuantizeCustomOp(%arg0: tensor<1x1x1x1xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "custom_op"}} {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
  %w_1 = arith.constant dense<127.0> : tensor<4096x1x1x1xf32>

  return %2 : tensor<?x2x2x1xi8>
}

// -----

// CHECK-LABEL: func @convolution_add_add_f32
func.func @convolution_add_add_f32(
    %lhs: tensor<?x3x2x1xf32>, %rhs: tensor<2x1x1x1xf32>,
    %zp_offset: tensor<?x2x2x1xf32>, %bias: tensor<1xf32>
  ) -> tensor<?x2x2x1xf32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[conv:.*]], %[[zp_offset:.*]]

// CustomOpNotWeightOnly-LABEL: QuantizeCustomOp
func.func @QuantizeCustomOp(%arg0: tensor<1x1x1x1xf32>) -> tensor<*xf32> attributes {tf.entry_function = {inputs = "input", outputs = "custom_op"}} {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
  %w = arith.constant dense<127.0> : tensor<1024x1x1x1xf32>

//     {0}, then it returns: tensor<1x3x5xi32>.
//   * If `xla_gather_op_output_type` == tensor<3x5xf32> and `collapsed_dims` ==
//     {1, 3}, then it returns: tensor<3x1x5x1xf32>.
Type GetSliceOpOutputType(Type xla_gather_op_output_type,
                          const absl::flat_hash_set<int64_t>& collapsed_dims) {
  if (auto ranked_output_type =