diff --git a/tools/mlir-opt/DefaultInitializeMLIRContext.cpp b/tools/mlir-opt/DefaultInitializeMLIRContext.cpp
new file mode 100644
index 0000000..df59132
--- /dev/null
+++ b/tools/mlir-opt/DefaultInitializeMLIRContext.cpp
@@ -0,0 +1,9 @@
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/TensorFlow/ControlFlowOps.h"
+
+using namespace mlir;
+
+// Register the TFControlFlow ops with the MLIRContext.
+void initializeMLIRContext(MLIRContext *ctx) {
+  TFControlFlow::registerOperations(*ctx);
+}
diff --git a/tools/mlir-opt/mlir-opt.cpp b/tools/mlir-opt/mlir-opt.cpp
index 6eb31d6..29fc9c4 100644
--- a/tools/mlir-opt/mlir-opt.cpp
+++ b/tools/mlir-opt/mlir-opt.cpp
@@ -85,9 +85,10 @@
   return result;
 }
 
-static void initializeMLIRContext(MLIRContext &ctx) {
-  TFControlFlow::registerOperations(ctx);
-}
+// The function to initialize the MLIRContext for different ops is defined in
+// another compilation unit to allow different tests to link in different
+// context initializations (e.g., op registrations).
+extern void initializeMLIRContext(MLIRContext *ctx);
 
 /// Parses the memory buffer and, if successfully parsed, prints the parsed
 /// output. Optionally, convert ML functions into CFG functions.
@@ -99,7 +100,7 @@
 
   // Parse the input file.
   MLIRContext context;
-  initializeMLIRContext(context);
+  initializeMLIRContext(&context);
   std::unique_ptr<Module> module(parseSourceFile(sourceMgr, &context));
   if (!module)
     return OptFailure;
@@ -219,7 +220,7 @@
 
     // Parse the input file.
     MLIRContext context;
-    initializeMLIRContext(context);
+    initializeMLIRContext(&context);
 
     // TODO: refactor into initializeMLIRContext so the normal parser pass
     // gets to use this.