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

[stellaraccident]

Summary

• This was left out of the previous commit accidentally.

[ftynse]

mlir/test/python/dialects/sparse_tensor/dialect.py

3–4

Is this TODO still relevant?

[stellaraccident]

Comments.

mlir/python/mlir/_cext_loader.py

def _reexport_cext(cext_module_name, target_module_name):
 
   Typically:
     from ._cext_loader import _reexport_cext
     _reexport_cext("ir", __name__)
     del _reexport_cext
   """
   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))
 
 
 # Add our 'dialects' parent module to the search path for implementations.
 _cext.globals.append_dialect_search_prefix("mlir.dialects")

mlir/python/mlir/dialects/sparse_tensor.py

+#  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

mlir/test/python/dialects/sparse_tensor/dialect.py

 # 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__)
   f()
   return f
 
 
 # CHECK-LABEL: TEST: testEncodingAttr1D