1 files changed, 9 insertions, 4 deletions
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index 397db4c20e74..da8c2167dc36 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 using namespace mlir::python::adaptors;
@@ -21,12 +22,16 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
   auto standaloneM = m.def_submodule("standalone");
 
   standaloneM.def(
-      "register_dialect",
+      "register_dialects",
       [](MlirContext context, bool load) {
-        MlirDialectHandle handle = mlirGetDialectHandle__standalone__();
-        mlirDialectHandleRegisterDialect(handle, context);
+        MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__();
+        MlirDialectHandle standaloneHandle =
+            mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(arithHandle, context);
+        mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(arithHandle, context);
+          mlirDialectHandleRegisterDialect(standaloneHandle, context);
         }
       },
       py::arg("context") = py::none(), py::arg("load") = true);