.lookupSymbol<mlir::func::FuncOp>(cluster_func.getFunc());

  bool use_spmd = false;
  if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>(kUseSpmdAttr))
    use_spmd = use_spmd_attr.getValue();

  auto num_cores_per_replica_attr =
      cluster_func->getAttrOfType<IntegerAttr>(kNumCoresPerReplicaAttr);
  if (!num_cores_per_replica_attr)
    return cluster_func.emitOpError(

  // TODO - b/310539922: Remove reference to TF/TFL utils.
  if (call_op->hasAttr(kQuantTraitAttrName)) {
    return;
  }

  auto function_name = call_op
                           ->getAttrOfType<FlatSymbolRefAttr>(
                               TF::kStablehloEntryFunctionAttrName)
                           .getValue();
  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

bool HasOutsideCompilationAttribute(Operation* op) {
  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {

  // mapping back from the inputs to the input arguments.

  llvm::SmallString<64> op_name;
  if (IsLegacyCallInstruction(inst)) {
    // The op_name is the name of the function.
    op_name.append(inst->getAttrOfType<mlir::SymbolRefAttr>("f")
                       .getLeafReference()
                       .getValue());
    // Remove the attribute from the instruction as it is already converted to
    // op_name.

}

mlir::LogicalResult GetTopology(mlir::tf_device::ClusterOp cluster,
                                std::string& topology) {
  mlir::StringAttr topology_attr =
      cluster->getAttrOfType<mlir::StringAttr>(tensorflow::kTopologyAttr);
  if (topology_attr) {
    topology = topology_attr.getValue();
    return mlir::success();
  } else {
    return cluster.emitOpError(

LogicalResult AssignDevicesInRegion(const Dialect* tf_dialect,
                                    tf_device::LaunchOp launch,
                                    Region& region) {
  auto parallel_group_attr =
      launch->getAttrOfType<StringAttr>(TF::kParallelExecAnnotation);
  auto result = region.walk([&](Operation* op) -> WalkResult {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();

    if (parallel_group_attr) {

    auto bias =
        rewriter.create<TF::ConstOp>(op->getLoc(), bias_type, bias_attr);

    if (op->getAttrOfType<StringAttr>("padding").getValue() == "EXPLICIT") {
      // Add Const op for padding value.
      ArrayRef<Attribute> padding_attr_array =
          op->getAttrOfType<ArrayAttr>("explicit_paddings").getValue();

      auto get_int = [](Attribute attr) {

// "arg0:0,arg1:0"}`, then this function returns `{"arg0:0", "arg1:0"}`.
SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;
  mlir::dyn_cast_or_null<StringAttr>(entry_function_attr.get("inputs"))
      .strref()
      .split(inputs, /*Separator=*/",");

      DeviceNameUtils::AddressSpace(parsed_name));
  return result.empty() ? kLocalhost : result;
}

std::string GetHost(Operation *op) {
  std::string device = "";
  if (StringAttr attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
    device = attr.getValue().str();
  }
  return GetHost(device);
}

// The device is considered to be on the localhost iff one of the following is
// true:

  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());
    // Skip non-tpu device cluster_func.
    auto cluster_id =
        cluster_func_op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    llvm::SmallVector<int, 4> dynamic_shape_arg_index;

    // Traverse the operands of the cluster func op and find which operand