#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kReplicaIdAttr[] = "_xla_replica_id";
constexpr char kDeviceOrdinalAttr[] = "device_ordinal";

#define GEN_PASS_DEF_REPLICATETOISLANDPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

namespace tf {
namespace libtf {
namespace runtime {

using ::tensorflow::testing::StatusIs;
using ::testing::HasSubstr;
using ::tf::libtf::impl::TaggedValueTensor;

constexpr char kSimpleModel[] =
    "tensorflow/cc/experimental/libtf/tests/testdata/simple-model";

TEST_P(RuntimeTest, SimpleModelCallableFloatTest) {
  Runtime runtime = RuntimeTest::GetParam()();

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<1xi32> {

};

TEST_F(TestPassInstrumentation, CreatedCalledAndSetsPassName) {
  RegisterPassInstrumentor(kTestInstrumentationName, [&]() {
    return std::make_unique<TestInstrumentor>(this);
  });
  constexpr char legalization[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

namespace {

using tensorflow::experimental::cc::Runtime;
using tensorflow::experimental::cc::RuntimeBuilder;
using tensorflow::experimental::cc::SavedModelAPI;
using tensorflow::experimental::cc::Status;

constexpr char kTestData[] = "cc/saved_model/testdata";

std::string SavedModelPath(tensorflow::StringPiece saved_model_dir) {
  return tensorflow::io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),

#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

#define DEBUG_TYPE "tf-executor-tpu-v1-island-inlining"

namespace mlir {
namespace tf_executor {

namespace {
constexpr llvm::StringRef kNestedModule = "_tpu_v1_compat_outlined";

#define GEN_PASS_DEF_EXECUTORTPUV1ISLANDINLININGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

using mlir::ValueRange;
using mlir::WalkResult;

constexpr llvm::StringRef kDeviceAttr = "device";
constexpr llvm::StringRef kNameAttr = "name";
constexpr llvm::StringRef kNumCoresPerReplicaAttr = "num_cores_per_replica";
constexpr llvm::StringRef kNumReplicasAttr = "num_replicas";
constexpr llvm::StringRef kMirroredVariableIndicesAttr =
    "_mirrored_variable_indices";

constexpr llvm::StringRef kBadReplicateInfoAttrMsg =

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

constexpr char kEntryFuncName[] = "main";

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {

      DenseFPElementsAttr value_attr =
          mlir::cast<DenseFPElementsAttr>(op.getValue());
      new_values.reserve(value_attr.getNumElements());

      constexpr float kMaxFloat16Value = 65504.f;
      constexpr float kMinFloat16Value = -65504.f;

      for (auto value : value_attr.template getValues<float>()) {
        new_values.push_back(Eigen::half(

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";

// Checks if a function only contains a tf_executor.graph.
bool IsSupportedGraph(func::FuncOp func) {
  if (!llvm::hasSingleElement(func)) return false;