diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 41435add404c..5b877aeb6858 100644
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -61,6 +61,18 @@ public:
   /// Return the type of this function.
   FunctionType getType() { return type; }
 
+  /// Change the type of this function in place. This is an extremely dangerous
+  /// operation and it is up to the caller to ensure that this is legal for this
+  /// function, and to restore invariants:
+  ///   - the entry block args must be updated to match the function params.
+  ///  - the arguments attributes may need an update: if the new type has less
+  ///    parameters we drop the extra attributes, if there are more parameters
+  ///    they won't have any attributes.
+  void setType(FunctionType newType) {
+    type = newType;
+    argAttrs.resize(type.getNumInputs());
+  }
+
   MLIRContext *getContext();
   Module *getModule() { return module; }