diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
--- a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
@@ -197,7 +197,7 @@
     f64 = ir.F64Type.get()
     # Be careful about increasing this because
     #     len(types) = 1 + len(level_choices)^rank * rank! * len(bitwidths)^2
-    shape = range(2, 6)
+    shape = range(2, 3)
     rank = len(shape)
     # All combinations.
     # TODO: add singleton here too; which requires updating how `np_arg0`