m.walk([&](tf_device::ClusterFuncOp cluster_func) {
    auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>();
    if (!replicate) return;
    auto mirrored_variable_indices_attr =
        replicate->getAttrOfType<ArrayAttr>(kMirroredVariableIndicesAttr);
    llvm::SmallDenseSet<int64_t, 8> mirrored_replicate_args;
    if (mirrored_variable_indices_attr) {
      for (const auto& mirrored_index : mirrored_variable_indices_attr) {

}

// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.
StringRef GetSharedName(Operation* op) {
  if (!op->hasAttrOfType<StringAttr>(kSharedNameAttr)) return "";
  return op->getAttrOfType<StringAttr>(kSharedNameAttr).getValue();
}

// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
// TODO(b/284222084): Move executor dialect utilities to a new library.

  // 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) {

      module_op, /*initializer_type=*/kTfSavedModelInitializerInitType);

  EXPECT_THAT(init_func_op, NotNull());
  EXPECT_THAT(init_func_op.getSymName(), "init_func");
  EXPECT_THAT(
      init_func_op->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr),
      kTfSavedModelInitializerInitType);
}

TEST_F(TfSavedModelTest, GetInitializerFunctionNoMatchingInitializerType) {
  constexpr StringRef kModuleOpStr = R"mlir(

                      ArrayRef<const Device *> results) override {
    if (llvm::isa<TF::TPUExecuteOp>(op) ||
        llvm::isa<TF::TPUExecuteAndUpdateVariablesOp>(op)) {
      auto device = op->getAttrOfType<StringAttr>(kDevice);
      for (auto *operand : operands)
        propagateIfChanged(operand, operand->SetDevice(device));
    } else {
      // Propagate device through other ops. These ops might have their

// "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=*/",");

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) {

  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

  auto clean_status = MakeCleanup([status] { TF_DeleteStatus(status); });

  // Builds TF operation and sets all the attributes.
  std::string node_name = "unnamed";
  if (auto attr = inst->getAttrOfType<mlir::StringAttr>("name")) {
    node_name = std::string(attr.getValue());
  }
  auto node_def_or = ConvertTFDialectOpToNodeDef(
      inst, node_name.c_str(), /*ignore_unregistered_attrs=*/true);