// the root cluster's host device. `callback` returning WasInterrupted
// indicates failure.
// The host device is null when the tpu_cluster HasModelParallelism: The
// HasModelParallelism case is currently unsupported in combination with
// outside compilation.
mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module, std::function<WalkResult(Operation*, tf_device::ClusterOp,

std::string GetDeviceAliasForHostOfLogicalCore(int core_index);

// Returns true if cluster contains model parallelism based on
// `num_cores_per_replica_attribute`. Otherwise returns false.
bool HasModelParallelism(mlir::tf_device::ClusterOp cluster);

// Returns true if the devices list contain any TPU devices
bool HasTPUDevice(const mlir::TF::RuntimeDevices& devices);

  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

  EXPECT_FALSE(HasModelParallelism(cluster));
}

TEST(TPURewriteDeviceUtilTest, TestHasModelParallelismTrue) {
  mlir::MLIRContext context;
  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();

}

std::string GetDeviceAliasForHostOfLogicalCore(const int core_index) {
  return llvm::formatv("{0}_{1}", kTPUReplicatedHost, core_index).str();
}

bool HasModelParallelism(mlir::tf_device::ClusterOp cluster) {
  mlir::IntegerAttr num_cores_per_replica_attr =
      cluster->getAttrOfType<mlir::IntegerAttr>(
          tensorflow::kNumCoresPerReplicaAttr);