diff --git a/include/mlir/IR/StandardOps.h b/include/mlir/IR/StandardOps.h
index 4c71b7e..99cbc59 100644
--- a/include/mlir/IR/StandardOps.h
+++ b/include/mlir/IR/StandardOps.h
@@ -299,7 +299,6 @@
   explicit StoreOp(const Operation *state) : OpBase(state) {}
 };
 
-/// TODO: change comment.
 /// The "return" operation represents a return statement of an ML function.
 /// The operation takes variable number of operands and produces no results.
 /// The operand number and types must match the signature of the ML function
diff --git a/lib/IR/StandardOps.cpp b/lib/IR/StandardOps.cpp
index c3f815f..e78290e 100644
--- a/lib/IR/StandardOps.cpp
+++ b/lib/IR/StandardOps.cpp
@@ -369,7 +369,6 @@
 bool ReturnOp::parse(OpAsmParser *parser, OperationState *result) {
   SmallVector<OpAsmParser::OperandType, 2> opInfo;
   SmallVector<Type *, 2> types;
-  SmallVector<SSAValue *, 2> operands;
 
   return parser->parseOperandList(opInfo, -1, OpAsmParser::Delimiter::None) ||
          (!opInfo.empty() && parser->parseColonTypeList(types)) ||
@@ -391,17 +390,15 @@
 const char *ReturnOp::verify() const {
   // ReturnOp must be part of an ML function.
   if (auto *stmt = dyn_cast<OperationStmt>(getOperation())) {
-    StmtBlock *block = stmt->getBlock();
-
-    if (!block || !isa<MLFunction>(block) ||
-        &cast<MLFunction>(block)->back() != stmt)
+    MLFunction *func = dyn_cast_or_null<MLFunction>(stmt->getBlock());
+    if (!func || &func->back() != stmt)
       return "must be the last statement in the ML function";
 
     // Return success. Checking that operand types match those in the function
     // signature is performed in the ML function verifier.
     return nullptr;
   }
-  return "cannot occur in a CFG function.";
+  return "cannot occur in a CFG function";
 }
 
 //===----------------------------------------------------------------------===//
@@ -470,6 +467,6 @@
 /// Install the standard operations in the specified operation set.
 void mlir::registerStandardOperations(OperationSet &opSet) {
   opSet.addOperations<AddFOp, AffineApplyOp, AllocOp, ConstantOp, DimOp, LoadOp,
-                      StoreOp, ReturnOp>(
+                      ReturnOp, StoreOp>(
       /*prefix=*/"");
 }