}

bool TypeValidForXLA(const mlir::Type& type) {
  const mlir::Type elem = getElementTypeOrSelf(type);
  return !mlir::isa<mlir::TF::ResourceType>(elem) &&
         !mlir::isa<mlir::TF::StringType>(elem);
}

mlir::LogicalResult GetDeviceToHostMap(
    mlir::tf_device::ClusterOp cluster,
    llvm::SmallVector<std::string, 8>& core_to_host) {
  mlir::TF::RuntimeDevices devices;

        "//tensorflow/compiler/mlir/tensorflow",
        "//tensorflow/compiler/mlir/tensorflow:attribute_utils",
        "//tensorflow/compiler/mlir/tensorflow:convert_attr",
        "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
        "//tensorflow/compiler/mlir/tensorflow:convert_type",
        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
        "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Support/FileUtilities.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/cc/saved_model/loader.h"

#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/OperationSupport.h"  // from @llvm-project

        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "//tensorflow/compiler/mlir/lite:tf_tfl_passes",
        "//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
        "//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:target_hardware",
        "//tensorflow/compiler/mlir/tensorflow:cluster_util",
        "//tensorflow/compiler/mlir/tensorflow:error_util",

#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/OperationSupport.h"  // from @llvm-project
#include "mlir/IR/PatternMatch.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project

  size_t numRead;
  mlir::Attribute attr = mlir::parseAttribute(input, context, {}, &numRead);
  if (!attr || !mlir::isa<mlir::StringAttr>(attr)) {
    LOG(ERROR) << "Input is not parsable as a MLIR StringAttr.";
    return nullptr;
  }
  auto str_attr = mlir::cast<mlir::StringAttr>(attr);

  mlir::DialectRegistry registry;
  RegisterMlirInputDialects(registry);

#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/Matchers.h"  // from @llvm-project
#include "mlir/IR/OpDefinition.h"  // from @llvm-project

#include "tensorflow/core/util/device_name_utils.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using mlir::Block;
using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::NamedAttribute;
using mlir::NamedAttrList;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OpResult;
using mlir::Region;

TEST(TPURewriteDeviceUtilTest, TestHasModelParallelismFalse) {
  mlir::MLIRContext context;
  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module_ref->getBodyRegion());

  llvm::SmallVector<mlir::Type, 8> result_types;
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(