func.return
}

// CHECK-LABEL: func @variant
func.func @variant(%arg: tensor<!tf_type.variant>) {
  // CHECK: tf.ZerosLike
  // CHECK-SAME: T = !corert.variant
  %0 = "tf.ZerosLike"(%arg) {device = "/device:CPU:0", T = !tf_type.variant} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
  func.return
}

// Checks that TF quantized attrs are lowered to the corert types

        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-LABEL: func_with_tfl_attrs
func.func @func_with_tfl_attrs(%arg0: tensor<!tf_type.variant<tensor<2xi32>>>, %arg1: tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<2xi32>>> attributes {tf.entry_function = {inputs = "arg0,arg1", outputs = "arg0"}} {
  return %arg0 : tensor<!tf_type.variant<tensor<2xi32>>>
}

// -----

// -----

// 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)

  // CHECK: "tf.Iterator"
  // CHECK-SAME: 2x42
  %1, %2 = "tf.AnonymousIteratorV2"() {output_shapes = [
    #tf_type.shape<2x42>], output_types = [f32]} : () -> (tensor<!tf_type.resource>, tensor<!tf_type.variant>)
  // CHECK: "tf.Iterator"
  // CHECK-SAME: 3x42
  %3 = "tf.AnonymousIteratorV3"() {output_shapes = [
    #tf_type.shape<3x42>], output_types = [f32]} : () -> tensor<!tf_type.resource>
  return

  func.return %0 : tensor<!tf_type.variant<tensor<2xi32>>>
}

//===----------------------------------------------------------------------===//
// Matmul op legalizations.

      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"() {