}
    }
  }

  llvm::SmallVector<DictionaryAttr, 8> tmp;
  for (const auto& res : result_attrs) {
    tmp.push_back(builder.getDictionaryAttr(res));
  }
  func.setAllResultAttrs(tmp);
}

#define GEN_PASS_DEF_TPUDEVICEPROPAGATIONPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUDevicePropagation

                                       arg_types, res_types)));

    // Update arg/result attributes.
    func.setAllArgAttrs(arg_attrs);
    func.setAllResultAttrs(res_attrs);

    rewriter.setInsertionPointToStart(func.addEntryBlock());
    // In TFE, the function body is inlined in a GraphOp. Create a GraphOp
    // instance and move the regions from GraphFuncOp to GraphOp.

  func.setAllArgAttrs(
      llvm::to_vector<4>(llvm::map_range(arg_attrs, [&](NamedAttrList& list) {
        return list.getDictionary(context_);
      })));
  func.setAllResultAttrs(
      llvm::to_vector<4>(llvm::map_range(ret_attrs, [&](NamedAttrList& list) {
        return list.getDictionary(context_);
      })));

  return absl::OkStatus();
}