#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/host_context/async_value.h"  // from @tf_runtime
#include "tfrt/host_context/chain.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime

  // 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 = ""} : () -> ()

        %4 = "tf.B"(%2) : (tensor<?xi32>) -> tensor<?xi32>
        tf_device.return %4 : tensor<?xi32>
      }) {} : () -> tensor<?xi32>

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

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

func.func @nested_entry_functions() attributes {tf.entry_function = {}} {
  tf_executor.graph {
     %control = tf_executor.island wraps "tf.StatefulPartitionedCall"() {config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @func} : () -> ()
     tf_executor.fetch
  }
  func.return
}

func.func @func() attributes {tf.entry_function = {}} {
  func.return
}

// tf_executor.island {_parallel_execution_ids=r0:0,p0:0} wraps "tf.OpC"
// tf_executor.island {_parallel_execution_ids=r0:0,p0:1} wraps "tf.OpD"
// tf_executor.island {_parallel_execution_ids=r0:1} wraps "tf.OpB"
// tf_executor.island {_parallel_execution_ids=r0:1,p0:0} wraps "tf.OpC"
// tf_executor.island {_parallel_execution_ids=r0:1,p0:1} wraps "tf.OpD"
//
// "tf.OpA" will not have `_parallel_execution_ids` attr,

    %0:6 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
      %outputs_0, %control_1 = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) {device = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>

// tf_executor dialect.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ConvertGraphdefToMlir(
    const GraphDef& graphdef, const GraphDebugInfo& debug_info,
    const GraphImportConfig& specs, mlir::MLIRContext* context);

// Given a Graph, returns a MLIR module containing the graph, expressed with
// tf_executor dialect.

    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.VariableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shape = #tf_type.shape<>, shared_name = "v_load_44"} : () -> tensor<!tf_type.f32ref>

#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"

    // violate the invariant imposed by the GraphExport pipeline that every
    // IslandOp perfectly wraps a single op.
    auto control_type =
        mlir::tf_executor::ControlType::get(tf_dialect->getContext());
    getOperation().walk(
        [&control_type](mlir::tf_executor::IslandOp curr_island) {
          mlir::TF::SplitIsland(curr_island, control_type);
        });
  }
}

}  // anonymous namespace