// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32> loc("RecvTPUEmbedding")

tf_executor::IslandOp CreateIsland(TypeRange result_types,
                                   const tf_executor::ControlType& control_type,
                                   const Location& loc, Operation& sub_op,
                                   tf_executor::IslandOp original_island) {
  OpBuilder builder(original_island);
  auto island = builder.create<tf_executor::IslandOp>(

#include "tensorflow/core/common_runtime/eager/eager_executor.h"

struct TFE_Executor {
  explicit TFE_Executor(bool async, bool enable_streaming_enqueue,
                        int in_flight_nodes_limit)
      : owned_executor(new tensorflow::EagerExecutor(
            async, enable_streaming_enqueue, in_flight_nodes_limit)) {}

  explicit TFE_Executor(tensorflow::EagerExecutor* executor)

  let summary = "Lifts tf_executor.island inner ops from a tf_executor.graph";

  let description = [{
    This pass converts tf_executor.graphs consisting of only tf_executor.islands and
    a tf_executor.fetch into a sea of nodes consisting of TensorFlow Dialect ops by
    lifting such ops out of a tf_executor.graph's tf_executor.islands. If V1 control

  %graph = tf_executor.graph {
    // This node would be renamed to bar1 [note: if imported from TF graphdef this would not be possible]
    %2:2 = tf_executor.island wraps "tf.Identity"(%arg1) {device = "", dtype = "tfdtype$DT_INT32"} : (tensor<10xi32>) -> tensor<10xi32> loc ("bar")
    // The following node would be renamed to bar2

  %0 = tf_executor.graph {
    %outputs_14, %control_15 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<1.0> : tensor<186xf32>} : () -> tensor<186xf32>
    %outputs_16, %control_17 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<2.0> : tensor<186xf32>} : () -> tensor<186xf32>

  func.func @main() {
    tf_executor.graph {
      %control = tf_executor.island wraps "tf.NoOp"() {_f = #tf_type.func<@callee, {attr2 = true, attr3 = 8.0 : f32}>} : () -> ()
      %control_1 = tf_executor.island(%control) wraps "tf.LegacyCall"() {f = @callee} : () -> ()
      tf_executor.fetch
    }
    func.return
  }

// CHECK: module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 919 : i32}}
module  {
  // CHECK: tf_executor.graph
  tfg.graph #tf_type.version<producer = 919, min_consumer = 12> {

func.func @controls_per_replica() {
  tf_executor.graph {
    %1 = tf_executor.ControlTrigger {}
    %2 = tf_executor.ControlTrigger {}
    %3 = tf_executor.island(%1, %2) {
      tf_device.replicate {n = 2 : i32} {
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK: %[[CT_0:.*]] = tf_executor.ControlTrigger

    tf_executor.fetch
  }
  func.return
}

// -----

// Check that we don't report a dependency between two stateful ops.
func.func @two_stateful_ops() -> () {
  tf_executor.graph {
    %island1 = tf_executor.island {
      "tf.OpA"() {is_stateless = false}: () -> ()
      tf_executor.yield
    }
    %island2 = tf_executor.island(%island1) {
      "tf.OpB"() {is_stateless = false}: () -> ()
      tf_executor.yield
    }