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

// CHECK: tfrt.return [[pred_res]]#0, [[body_res]]#1, [[pred_res]]#1 : !tfrt.chain, !tfrt_fallback.tf_tensor, i1

// CHECK: func @"while_cond_lt9/tfrt_predicate"([[ch:%.*]]: !tfrt.chain, [[arg:%.*]]: !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)

    visibility = [
        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
        "//tensorflow/core/runtime_fallback:internal",
        "//tensorflow/core/tfrt/saved_model:friends",
        # copybara:uncomment "//third_party/tf_runtime_google:__subpackages__",

#include "tfrt/host_context/function.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime
#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
#include "tfrt/support/ref_count.h"  // from @tf_runtime

namespace tensorflow {

using ::tfrt::AsyncValue;
using ::tfrt::BEFFile;
using ::tfrt::ExecutionContext;
using ::tfrt::Function;
using ::tfrt::MakeAvailableAsyncValueRef;
using ::tfrt::RCReference;

#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime

        "//tensorflow/core/tfrt/ifrt/...",
        "//tensorflow/core/tfrt/kernels/...",
        "//tensorflow/core/tfrt/mlrt/...",
        "//tensorflow/core/tfrt/saved_model/...",
        "//tensorflow/core/tfrt/tfrt_session/...",
    ] + if_google([
        "//learning/brain/experimental/mlir/tflite/tfmrt/...",
        "//learning/brain/experimental/mlir/tfrt_compiler/...",
        "//learning/brain/experimental/tfrt/...",

#include "tensorflow/core/platform/threadpool_interface.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime
#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
#include "tfrt/support/ref_count.h"  // from @tf_runtime

namespace tensorflow {

  %s = "tfrt_test.get_string"() { value = "Running @matmul_cpu" } : () -> !tfrt.string
  %ch4 = "tfrt_test.print_string"(%s, %ch3) : (!tfrt.string, !tfrt.chain) -> (!tfrt.chain)

  %ch5 = "tfrt_fallback_async.print_tensor"(%result, %ch4) : (!tfrt_fallback.tf_tensor, !tfrt.chain) -> !tfrt.chain

  tfrt.return %ch5, %result : !tfrt.chain, !tfrt_fallback.tf_tensor
}

#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "tfrt/core_runtime/opdefs/traits.h"  // from @tf_runtime
#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime

using namespace mlir;  // NOLINT

namespace tfrt {
namespace fallback_sync {

// Dialect for fallback operations.
class FallbackSyncDialect : public Dialect {
 public:

#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime

namespace tfrt {
namespace fallback_async {
namespace {

TEST(SavedModelTest, MapFallbackArgs) {
  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir");

  mlir::DialectRegistry registry;
  RegisterTFRTDialects(registry);