output_types = [f32],
        metadata = ""
      } :
         (tensor<!tf_type.string>, tensor<*x!tf_type.string>, tensor<!tf_type.string>, tensor<*xi64>, tensor<!tf_type.string>) -> tensor<*x!tf_type.variant>
      %add1 = "tf.Add"(%str, %arg3) : (tensor<!tf_type.string>, tensor<!tf_type.string>) -> tensor<!tf_type.string>
      // CHECK: tf_executor.island(%[[CONTROL]]) wraps "tf.GeneratorDataset"

      %outputs_12, %control_13 = tf_executor.island wraps "tf.Identity"(%outputs_10) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>

}

def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";

  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a bytecode function.

    Example:
      %res = tf_mlrt.batch_function(%input, %captured_input)  {
          device = "/device:CPU:0",

    : (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant>
  func.return %0 : tensor<*x!tf_type.variant>

// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() <{value = dense<1000> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() <{value = dense<8> : tensor<i64>}> : () -> tensor<i64>

// -----

// Check that merge data inputs of variant type are broadcastable to the output
func.func @invalid_merge(%arg0: tensor<*x!tf_type.variant>, %arg1: tensor<4x!tf_type.variant>) -> tensor<8x!tf_type.variant> {
  %result = tf_executor.graph {
    %value, %idx, %ctlMerge = "tf_executor.Merge"(%arg0, %arg1) : (tensor<*x!tf_type.variant>, tensor<4x!tf_type.variant>) -> (tensor<8x!tf_type.variant>, tensor<i32>, !tf_executor.control)

      output_types = [!tf_type.string],
      metadata = ""} : (
              tensor<4xf32>,
              tensor<3xf32>,
              tensor<!tf_type.resource>,
              tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

func.func @testIncompleteGeneratorDataset() {
  // expected-error@+1 {{'tf.GeneratorDataset' op failed to convert to region form}}
  %0 = "tf.GeneratorDataset"() {

    %0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
    return %0 : tensor<1x3x3x4xf32>
  }
// CHECK: @uniform_dequantize_0
}

// -----

// Tests a variant where there is no stablehlo.convert op in between the
// filter constant and the convolution op.
//
// `filter (f32) -> convolution`
//
// instead of:
//
// `filter (i8) -> convert (i8 -> f32) -> convolution`

// -----

func.func @testIfThen(%arg0: tensor<!tf_type.variant>) -> tensor<!tf_type.variant> {
  func.return %arg0 : tensor<!tf_type.variant>
}
func.func @testIfElse(%arg0: tensor<!tf_type.variant>) -> tensor<!tf_type.variant> {
  func.return %arg0 : tensor<!tf_type.variant>
}

// CHECK-LABEL: func @testIfCasts(%arg0: tensor<i1>, %arg1: tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>

    %2 = "tf.Identity"(%1) : (tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
    "tf.Identity"(%2) : (tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
    tf_device.return
  }) {cluster_attr = "cluster_attr"} : () -> ()
  func.return
}

  return %result : tensor<i32>
}

// CHECK-LABEL: @NopMapFnBody
func.func private @NopMapFnBody(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>) -> () {