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

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

  // Following are helper functions to assist integrating TFRT with current
  // TF eager runtime.
  // TODO(b/172877902): These helper functions are currently used to support
  // PyFuncOp on TFRT, and might be useful for ops that directly use low
  // level TF APIs. Remove/replace the following functions when TFRT native
  // ops are implemented.
  //===--------------------------------------------------------------------===//

    cc_api_version = 2,
    visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
    ],
)

cc_library(
    name = "cluster_tf",
    srcs = ["cluster_tf.cc"],
    hdrs = ["cluster_tf.h"],
    visibility = [
        "//learning/serving/contrib/tfrt/mlir/saved_model_analysis:__pkg__",
        "//tensorflow/compiler/mlir/tensorflow/transforms:__pkg__",

  return tensorflow::io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
                                  kTestData, saved_model_dir);
}

// This value parameterized test allows us to test both TFRT
// and non TFRT runtimes.
// https://github.com/google/googletest/blob/dcc92d0ab6c4ce022162a23566d44f673251eee4/googletest/docs/advanced.md#value-parameterized-tests
class CSavedModelAPITest : public ::testing::TestWithParam<bool> {};

#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/global_state.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tensorflow/core/tpu/tpu_defs.h"
#include "tsl/framework/device_type.h"

namespace tensorflow {
namespace {

#ifdef TFRT_FALLBACK_ASYNC_OPS
#else
#define TFRT_FALLBACK_ASYNC_OPS

include "tfrt/tfrt_op_base.td"
include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"
include "tfrt/compiler/opdefs/tfrt_traits.td"
include "tfrt/core_runtime/opdefs/corert_base.td"
include "tfrt/core_runtime/opdefs/corert_traits.td"
include "tfrt_fallback.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

func.func @f(%arg0: !tfrt.chain, %arg1: !tfrt_fallback.tf_tensor, %arg2: !tfrt_fallback.tf_tensor, %arg3: !tfrt_fallback.tf_tensor, %arg4: !tfrt_fallback.tf_tensor,
        %arg5: !tfrt_fallback.tf_tensor, %arg6: !tfrt_fallback.tf_tensor, %arg7: !tfrt_fallback.tf_tensor, %arg8: !tfrt_fallback.tf_tensor, %arg9: !tfrt_fallback.tf_tensor,

#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tensorflow/core/tpu/tpu_defs.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;