tf_device.return %5: tensor<i64>
          }) {device = "/TPU:2"} : () -> (tensor<i64>)
          tf_device.return %4 : tensor<i64>
        }) : () -> (tensor<i32>, tensor<i64>)
        tf_device.return %3#0, %3#1 : tensor<i32>, tensor<i64>
      }
      tf_executor.yield %2#0, %2#1, %2#2, %2#3 : tensor<i32>, tensor<i32>, tensor<i64>, tensor<i64>
    }

  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64
  // CHECK-SAME: num_batch_threads = 1 : i64
  // CHECK-SAME: shared_name = "batch/"

func.func @typeNotSupportedNotLegalized(%arg0: tensor<!tf_type.variant<tensor<*xf64>>>, %arg1: tensor<i32>, %arg2: tensor<*xf64>) -> tensor<!tf_type.variant<tensor<*xf64>>> {
  %0 = "tf.TensorListSetItem"(%arg0, %arg1, %arg2) : (tensor<!tf_type.variant<tensor<*xf64>>>, tensor<i32>, tensor<*xf64>) -> tensor<!tf_type.variant<tensor<*xf64>>>
  // CHECK-NOT: "tfl.custom"

  %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
  } : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
  %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
  } : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>

  %1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>...

"tf.BatchFunction"(%cst) <{allowed_batch_sizes = [1], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", enable_large_batch_splitting = true, f = @_batched, low_priority_allowed_batch_sizes = [], low_priority_batch_timeout_micros = 0 : i64, low_priority_max_batch_size = 0 : i64, low_priority_max_enqueued_batches = 0 : i64, max_batch_size = 1 : i64, max_enqueued_batches = 1 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch_function___inf...

// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
// CHECK-NEXT:    %2 = "tf.IfrtCall"(%1, %arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []

      dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
      window = {stride = [1, 1], pad = [[0, 0], [0, 0]],
      lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
      {batch_group_count = 1 : i64, feature_group_count = 1 : i64}
    : (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>

  %filter = mhlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>

    // expected-error @+1 {{'tf.TPUPartitionedInput' op TF2XLA TPU bridge input check: number of inputs inconsistent. num_cores_per_replica=2 no. of inputs=3}}
    %pi, %c0 = tf_executor.island wraps "tf.TPUPartitionedInput"(%arg0, %arg1, %arg1) {index = 1 : i64} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>