Add information displayed on failure, renamed variables

Author: oleksandr-pavlyk

dpctl/tests/test_usm_ndarray_top_k.py

@@ -45,16 +45,16 @@ def test_topk_1d_largest(dtype, n):
 
     o = dpt.ones(n, dtype=dtype)
     z = dpt.zeros(n, dtype=dtype)
-    zo = dpt.concat((o, z))
-    inp = dpt.roll(zo, 734)
+    oz = dpt.concat((o, z))
+    inp = dpt.roll(oz, 734)
     k = 5
 
     s = dpt.top_k(inp, k, mode="largest")
     assert s.values.shape == (k,)
     assert s.values.dtype == inp.dtype
     assert s.indices.shape == (k,)
-    assert dpt.all(s.values == dpt.ones(k, dtype=dtype))
-    assert dpt.all(s.values == inp[s.indices])
+    assert dpt.all(s.values == dpt.ones(k, dtype=dtype)), s.values
+    assert dpt.all(s.values == inp[s.indices]), s.indices
 
 
 @pytest.mark.parametrize(
@@ -82,34 +82,38 @@ def test_topk_1d_smallest(dtype, n):
 
     o = dpt.ones(n, dtype=dtype)
     z = dpt.zeros(n, dtype=dtype)
-    zo = dpt.concat((o, z))
-    inp = dpt.roll(zo, 734)
+    oz = dpt.concat((o, z))
+    inp = dpt.roll(oz, 734)
     k = 5
 
     s = dpt.top_k(inp, k, mode="smallest")
     assert s.values.shape == (k,)
     assert s.values.dtype == inp.dtype
     assert s.indices.shape == (k,)
-    assert dpt.all(s.values == dpt.zeros(k, dtype=dtype))
-    assert dpt.all(s.values == inp[s.indices])
+    assert dpt.all(s.values == dpt.zeros(k, dtype=dtype)), s.values
+    assert dpt.all(s.values == inp[s.indices]), s.indices
 
 
 # triage failing top k radix implementation on CPU
 # replicates from Python behavior of radix sort topk implementation
 @pytest.mark.parametrize("n", [33, 255, 511, 1021, 8193])
-def test_topk_largest_1d_radix_i1_255(n):
+def test_topk_largest_1d_radix_i1(n):
     get_queue_or_skip()
     dt = "i1"
 
     o = dpt.ones(n, dtype=dt)
     z = dpt.zeros(n, dtype=dt)
-    zo = dpt.concat((o, z))
-    inp = dpt.roll(zo, 734)
+    oz = dpt.concat((o, z))
+    inp = dpt.roll(oz, 734)
     k = 5
 
-    sorted = dpt.copy(dpt.sort(inp, descending=True, kind="radixsort")[:k])
-    argsorted = dpt.copy(
-        dpt.argsort(inp, descending=True, kind="radixsort")[:k]
-    )
-    assert dpt.all(sorted == dpt.ones(k, dtype=dt))
-    assert dpt.all(sorted == inp[argsorted])
+    sorted_v = dpt.sort(inp, descending=True, kind="radixsort")
+    argsorted = dpt.argsort(inp, descending=True, kind="radixsort")
+
+    assert dpt.all(sorted_v == inp[argsorted])
+
+    topk_vals = dpt.copy(sorted_v[:k])
+    topk_inds = dpt.copy(argsorted[:k])
+
+    assert dpt.all(topk_vals == dpt.ones(k, dtype=dt))
+    assert dpt.all(topk_vals == inp[topk_inds]), topk_inds