explicit_package (#82)

Author: makslevental

setup.py

@@ -1,7 +1,7 @@
 import os
 
 from pip._internal.req import parse_requirements
-from setuptools import setup
+from setuptools import setup, find_namespace_packages
 
 # TODO: find from extras maybe
 HOST_MLIR_PYTHON_PACKAGE_PREFIX = os.environ.get(
@@ -15,6 +15,12 @@ def load_requirements(fname):
     return [str(ir.requirement) for ir in reqs]
 
 
+packages = [
+    f"{HOST_MLIR_PYTHON_PACKAGE_PREFIX}.extras.{p}"
+    for p in find_namespace_packages(where="mlir/extras")
+] + [f"{HOST_MLIR_PYTHON_PACKAGE_PREFIX}.extras"]
+
+
 setup(
     name=PACKAGE_NAME,
     version="0.0.7",
@@ -28,6 +34,7 @@ def load_requirements(fname):
         "mlir": ["mlir-python-bindings"],
     },
     python_requires=">=3.8",
+    packages=packages,
     # lhs is package namespace, rhs is path (relative to this setup.py)
     package_dir={
         f"{HOST_MLIR_PYTHON_PACKAGE_PREFIX}.extras": "mlir/extras",

tests/test_other_hosts.py

@@ -6,16 +6,6 @@
 @pytest.mark.skipif(jax_not_installed(), reason="jax not installed")
 def test_jax_trampolines_smoke():
     # noinspection PyUnresolvedReferences
-    from jaxlib.mlir.dialects import (
-        arith,
-        builtin,
-        chlo,
-        func,
-        math,
-        memref,
-        mhlo,
-        scf,
-        sparse_tensor,
-        stablehlo,
-        vector,
-    )
+    from jaxlib.mlir.extras import context
+    # noinspection PyUnresolvedReferences
+    from jaxlib.mlir.extras.runtime import passes