[Mlir-commits] [mlir] f38633d - [mlir][Python] Re-export cext sparse_tensor module to the public namespace.

[Stella Laurenzo]

Author: Stella Laurenzo
Date: 2021-05-10T18:08:29Z
New Revision: f38633d1bbf5842f37ad722a2f0edfdfd80733a2

URL: https://github.com/llvm/llvm-project/commit/f38633d1bbf5842f37ad722a2f0edfdfd80733a2
DIFF: https://github.com/llvm/llvm-project/commit/f38633d1bbf5842f37ad722a2f0edfdfd80733a2.diff

LOG: [mlir][Python] Re-export cext sparse_tensor module to the public namespace.

* This was left out of the previous commit accidentally.

Differential Revision: https://reviews.llvm.org/D102183

Added: 
    mlir/python/mlir/dialects/sparse_tensor.py

Modified: 
    mlir/python/mlir/_cext_loader.py
    mlir/test/python/dialects/sparse_tensor/dialect.py

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/_cext_loader.py b/mlir/python/mlir/_cext_loader.py
index 35847efa9e4a0..3a9cde380354c 100644
--- a/mlir/python/mlir/_cext_loader.py
+++ b/mlir/python/mlir/_cext_loader.py
@@ -45,7 +45,10 @@ def _reexport_cext(cext_module_name, target_module_name):
   """
   import sys
   target_module = sys.modules[target_module_name]
-  source_module = getattr(_cext, cext_module_name)
+  submodule_names = cext_module_name.split(".")
+  source_module = _cext
+  for submodule_name in submodule_names:
+    source_module = getattr(source_module, submodule_name)
   for attr_name in dir(source_module):
     if not attr_name.startswith("__"):
       setattr(target_module, attr_name, getattr(source_module, attr_name))

diff  --git a/mlir/python/mlir/dialects/sparse_tensor.py b/mlir/python/mlir/dialects/sparse_tensor.py
new file mode 100644
index 0000000000000..687d1e6cdbc75
--- /dev/null
+++ b/mlir/python/mlir/dialects/sparse_tensor.py
@@ -0,0 +1,7 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._cext_loader import _reexport_cext
+_reexport_cext("dialects.sparse_tensor", __name__)
+del _reexport_cext

diff  --git a/mlir/test/python/dialects/sparse_tensor/dialect.py b/mlir/test/python/dialects/sparse_tensor/dialect.py
index f10116de2033a..9b8c66c327458 100644
--- a/mlir/test/python/dialects/sparse_tensor/dialect.py
+++ b/mlir/test/python/dialects/sparse_tensor/dialect.py
@@ -1,8 +1,7 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-# TODO: Import this into the user-package vs the cext.
-from _mlir.dialects import sparse_tensor as st
+from mlir.dialects import sparse_tensor as st
 
 def run(f):
   print("\nTEST:", f.__name__)