diff --git a/README.md b/README.md
index 26eb050da..489e472ec 100644
--- a/README.md
+++ b/README.md
@@ -31,18 +31,6 @@ This will build the image for the appropriate python version then copy the wheel
 
 ### Build wheels in your environment
 
-#### Temporary MLIR issue
-
-Due to an issue with MLIR, you will need to manually add `__init__.py` files to the `mlir` python package after the build.
-
-```bash
-$ make python-bindings
-$ touch build/tools/zamalang/python_packages/zamalang_core/mlir/__init__.py
-$ touch build/tools/zamalang/python_packages/zamalang_core/mlir/dialects/__init__.py
-```
-
-#### Build wheel
-
 Building the wheels is actually simple.
 
 ```bash
diff --git a/compiler/setup.py b/compiler/setup.py
index 48eccdc9a..1e5855260 100644
--- a/compiler/setup.py
+++ b/compiler/setup.py
@@ -38,7 +38,11 @@ setuptools.setup(
     url="https://github.com/zama-ai/homomorphizer",
     packages=setuptools.find_packages(
         where="build/tools/zamalang/python_packages/zamalang_core",
-        include=["zamalang", "zamalang.*", "mlir", "mlir.*"],
+        include=["zamalang", "zamalang.*"],
+    )
+    + setuptools.find_namespace_packages(
+        where="build/tools/zamalang/python_packages/zamalang_core",
+        include=["mlir", "mlir.*"],
     ),
     package_dir={"": "build/tools/zamalang/python_packages/zamalang_core"},
     include_package_data=True,