diff --git a/mlir/test/python/dialects/sparse_tensor/test_SpMM.py b/mlir/test/python/dialects/sparse_tensor/test_SpMM.py
--- a/mlir/test/python/dialects/sparse_tensor/test_SpMM.py
+++ b/mlir/test/python/dialects/sparse_tensor/test_SpMM.py
@@ -3,16 +3,17 @@
 import os
 import ctypes
 import mlir.all_passes_registration
+
 import numpy as np
+import mlir.ir as ir
+import mlir.dialects.sparse_tensor as st
+import mlir.runtime as rt
 
-from mlir.dialects import builtin
-from mlir.dialects.linalg.opdsl.lang import *
-from mlir.dialects.sparse_tensor import *
-from mlir.execution_engine import *
-from mlir.ir import *
-from mlir.passmanager import *
-from mlir.runtime import *
+from mlir.dialects.builtin import FuncOp
+from mlir.execution_engine import ExecutionEngine
+from mlir.passmanager import PassManager
 
+from mlir.dialects.linalg.opdsl.lang import *
 
 def run(f):
   print('\nTEST:', f.__name__)
@@ -28,29 +29,29 @@
   C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]
 
 
-def build_SpMM(attr: EncodingAttr):
+def build_SpMM(attr: st.EncodingAttr):
   """Build SpMM kernel.
 
   This method generates a linalg op with for matrix multiplication using
   just the Python API. Effectively, a generic linalg op is constructed
   that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A.
   """
-  module = Module.create()
+  module = ir.Module.create()
   f64 = ir.F64Type.get()
-  a = RankedTensorType.get([3, 4], f64, attr)
-  b = RankedTensorType.get([4, 2], f64)
-  c = RankedTensorType.get([3, 2], f64)
+  a = ir.RankedTensorType.get([3, 4], f64, attr)
+  b = ir.RankedTensorType.get([4, 2], f64)
+  c = ir.RankedTensorType.get([3, 2], f64)
   arguments = [a, b, c]
-  with InsertionPoint(module.body):
+  with ir.InsertionPoint(module.body):
 
-    @builtin.FuncOp.from_py_func(*arguments)
+    @FuncOp.from_py_func(*arguments)
     def spMxM(*args):
       return matmul_dsl(args[0], args[1], outs=[args[2]])
 
   return module
 
 
-def boilerplate(attr: EncodingAttr):
+def boilerplate(attr: st.EncodingAttr):
   """Returns boilerplate main method.
 
   This method sets up a boilerplate main method that calls the generated
@@ -75,11 +76,12 @@
 """
 
 
-def build_compile_and_run_SpMM(attr: EncodingAttr, support_lib: str, compiler):
+def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
+                               compiler):
   # Build.
   module = build_SpMM(attr)
   func = str(module.operation.regions[0].blocks[0].operations[0].operation)
-  module = Module.parse(func + boilerplate(attr))
+  module = ir.Module.parse(func + boilerplate(attr))
   # Compile.
   compiler(module)
   execution_engine = ExecutionEngine(
@@ -90,11 +92,11 @@
   Cin = np.zeros((3, 2), np.double)
   Cout = np.zeros((3, 2), np.double)
   Cin_memref_ptr = ctypes.pointer(
-      ctypes.pointer(get_ranked_memref_descriptor(Cin)))
+      ctypes.pointer(rt.get_ranked_memref_descriptor(Cin)))
   Cout_memref_ptr = ctypes.pointer(
-      ctypes.pointer(get_ranked_memref_descriptor(Cout)))
+      ctypes.pointer(rt.get_ranked_memref_descriptor(Cout)))
   execution_engine.invoke('main', Cout_memref_ptr, Cin_memref_ptr)
-  Cresult = ranked_memref_to_numpy(Cout_memref_ptr[0])
+  Cresult = rt.ranked_memref_to_numpy(Cout_memref_ptr[0])
 
   # Sanity check on computed result.
   expected = [[12.3, 12.0], [0.0, 0.0], [16.5, 19.8]]
@@ -121,7 +123,7 @@
         f'convert-std-to-llvm')
     self.pipeline = pipeline
 
-  def __call__(self, module: Module):
+  def __call__(self, module: ir.Module):
     PassManager.parse(self.pipeline).run(module)
 
 
@@ -130,7 +132,7 @@
 @run
 def testSpMM():
   support_lib = os.getenv('SUPPORT_LIB')
-  with Context() as ctx, Location.unknown():
+  with ir.Context() as ctx, ir.Location.unknown():
     count = 0
     # Fixed compiler optimization strategy.
     # TODO: explore state space here too
@@ -144,20 +146,20 @@
     # Exhaustive loop over various ways to annotate a kernel with
     # a *single* sparse tensor. Even this subset already gives
     # quite a large state space!
-    levels = [[DimLevelType.dense, DimLevelType.dense],
-              [DimLevelType.dense, DimLevelType.compressed],
-              [DimLevelType.compressed, DimLevelType.dense],
-              [DimLevelType.compressed, DimLevelType.compressed]]
+    levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
+              [st.DimLevelType.dense, st.DimLevelType.compressed],
+              [st.DimLevelType.compressed, st.DimLevelType.dense],
+              [st.DimLevelType.compressed, st.DimLevelType.compressed]]
     orderings = [
-        AffineMap.get_permutation([0, 1]),
-        AffineMap.get_permutation([1, 0])
+        ir.AffineMap.get_permutation([0, 1]),
+        ir.AffineMap.get_permutation([1, 0])
     ]
     bitwidths = [0, 8, 32]
-    for levels in levels:
+    for level in levels:
       for ordering in orderings:
         for pwidth in bitwidths:
           for iwidth in bitwidths:
-            attr = EncodingAttr.get(levels, ordering, pwidth, iwidth)
+            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
             compiler = SparseCompiler(options=opt)
             build_compile_and_run_SpMM(attr, support_lib, compiler)
             count = count + 1