}

// -----

// Check that tf_executor.graph can't be nested directly in a tf_executor.graph.
func.func @nested_graph() {
  tf_executor.graph {
    tf_executor.graph {}
// expected-error@-1 {{'tf_executor.graph' op unallowed directly inside another tf_executor.graph}}
  }
  func.return
}

// -----

// Check that a tf_executor.fetch is terminating a tf_executor.graph (custom parser)

      %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor<?xf32>) -> tensor<?xf32>

      // CHECK: tf_executor.yield %[[CALL_OUTPUT]]
      tf_executor.yield %3 : tensor<?xf32>
    }
    tf_executor.fetch %1#0 : tensor<?xf32>
  }
  func.return %0 : tensor<?xf32>
}

func.func @_func(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  func.return %arg0 : tensor<?xf32>
}

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"

#define DEBUG_TYPE "tf-functional-to-executor"

namespace mlir {

namespace {
// This pass converts mlir functions consisting of mlir ops into a tf_executor
// dialect as a single island.
// Result like so:
//   func @my_fn(%argi...) -> (result_t) {
//     %results:[[n_args]] = tf_executor.graph {

        tf_device.return %4 : tensor<?xi32>
      }) {device = "/device:test_device:0"} : () -> tensor<?xi32>

      // CHECK: tf_executor.yield %[[LAUNCH_OUTPUT]]
      tf_executor.yield %3 : tensor<?xi32>
    }
    tf_executor.fetch %1#0 : tensor<?xi32>
  }
  func.return %0 : tensor<?xi32>
}

// CHECK: func private @[[LAUNCH]]

    %0 = tf_executor.graph {
      %1, %c0 = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor<i1>} : () -> tensor<i1>
      tf_executor.fetch %1 : tensor<i1>
    }
    return %0 : tensor<i1>
  }

  func.func private @while_body(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
    tf_executor.graph {

      %control_90 = tf_executor.island(%control_28, %control_87) wraps "tf.NoOp"() {device = ""} : () -> ()
      %outputs_91, %control_92 = tf_executor.island(%control_90) wraps "tf.Identity"(%outputs_88) {device = ""} : (tensor<i1>) -> tensor<*xi1>
      tf_executor.fetch %outputs_91, %control_28, %control_87 : tensor<*xi1>, !tf_executor.control, !tf_executor.control
    }
    return %0 : tensor<*xi1>
  }

    tf_executor.graph {
      %out, %ctl = tf_executor.island wraps "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
      %out_0, %ctl_0 = tf_executor.island wraps "tf.ReadVariableOp"(%out) : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<2xf32>

  // CHECK-NOT: _replication_info =
  %0 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.GuaranteeConst"(%arg1) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %outputs_0, %control_1 = tf_executor.island wraps "tf.GuaranteeConst"(%arg2) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %control_2 = tf_executor.island wraps "tf.NoOp"() {_pivot_for_cluster = "cluster", device = ""} : () -> ()

func.func @multiple_func_body_ops() {
  tf_executor.graph {
    tf_executor.island {
      "tf.NoOp"() : () -> ()
      tf_executor.yield
    }
    tf_executor.fetch
  }
  tf_executor.graph {
    tf_executor.island {
      "tf.NoOp"() : () -> ()
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

// -----

  func.func @main() {
    tf_executor.graph {
      %control = tf_executor.island {
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    tf_executor.graph {
      %control = tf_executor.island {
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }