if (it == buffer_to_size.end()) continue;
    output_buffer_to_size.emplace_back(retval.index(), new_outputs.size(),
                                       it->getSecond().fixed);
    new_outputs.push_back(it->getSecond().size);
  }
  OpBuilder(old_terminator)
      .create<TerminatorOp>(old_terminator->getLoc(), new_outputs);
  old_terminator->erase();
  return output_buffer_to_size;
}

    if (info1_it == infos1.end()) continue;
    auto& info = result[entry.getFirst()];
    info = entry.getSecond();
    if (info.updated) continue;
    if (info1_it->getSecond().used) {
      info.used = true;
      info.updated = info1_it->getSecond().updated;
      info.data_type = info1_it->getSecond().data_type;
    }
  }
  return result;
}

  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    auto it = data_var_to_size_var.find(while_op.getOperand(i));
    if (it == data_var_to_size_var.end()) continue;
    new_while_operands.push_back(it->getSecond());
  }
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getInputs(), new_while_operands,
      while_op->getAttrs());

    auto it = value_to_device.find(operand);
    if (it == value_to_device.end()) return llvm::StringRef();

    if (new_device.empty()) {
      new_device = it->getSecond();
      continue;
    }

    if (new_device != it->getSecond()) return llvm::StringRef();
  }

  return new_device;
}

// Propagates devices from function arguments.
void PropagateDevicesFromArguments(
    func::FuncOp func,

    }
  }
  control_predecessors_.clear();
  for (auto& entry : sorted_control_predecessors_) {
    llvm::sort(entry.getSecond(), [](Operation* a, Operation* b) {
      return a->isBeforeInBlock(b);
    });
  }
  for (auto& entry : sorted_control_successors_) {
    llvm::sort(entry.getSecond(), [](Operation* a, Operation* b) {
      return a->isBeforeInBlock(b);
    });
  }

  for (auto &kv : constraints_) walk(kv.getFirst(), kv.getSecond());
}

std::optional<ValueConstraint> ValuesConstraintSet::GetConstraint(
    Value value) const {
  auto it = constraints_.find(value);
  if (it == constraints_.end()) return std::nullopt;
  return it->getSecond();
}

bool ValuesConstraintSet::HasConstraint(Value value) const {

  auto it = resource_value_to_ids_.find(resource);
  assert(it != resource_value_to_ids_.end() && !it->getSecond().empty());
  // The set is sorted so we only need to check the first element since
  // kUnknownResourceId < 0.
  static_assert(kUnknownResourceId < 0,
                "kUnknownResourceId should be negative");
  return *it->getSecond().begin() == kUnknownResourceId;
}

const llvm::SmallSet<int64_t, 8>&

  SymbolTable symbol_table(module);
  for (auto& func : funcs_to_rewrite) {
    if (func.getSecond().original.isPublic() &&
        !func.getSecond().original.symbolKnownUseEmpty(module)) {
      auto clone = func.getSecond().original.clone();
      clone.setPrivate();
      symbol_table.insert(clone);
      func.getSecond().clone = clone;
    }
  }

  return success();
}

    auto num_replicas = cluster_metadata->getSecond().get(kNumReplicasAttr);
    if (!num_replicas || !num_replicas.isa<mlir::IntegerAttr>())
      return cluster.emitError()
             << "requires '" << kNumReplicasAttr << "' int attribute";

    int num_cores_per_replica = 1;
    auto num_cores_per_replica_attr = mlir::dyn_cast_or_null<mlir::IntegerAttr>(
        cluster_metadata->getSecond().get(kNumCoresPerReplicaAttr));

  // placed in logical core 0.
  // TODO(b/148913020): Remove this constraint once model parallelism is
  // supported.
  assert(devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))
             ->getSecond()
             .size() == num_replicas);

  llvm::SmallVector<std::pair<ValueRange, Type>, 8> new_replicated_inputs;
  llvm::SmallVector<Value, 8> new_packed_inputs;