From a4d3eeb4f6cbbe78c7921d6f90fc2e903521c1b1 Mon Sep 17 00:00:00 2001 From: drisspg Date: Fri, 10 Jan 2025 09:15:35 -0800 Subject: [PATCH] Make it easer to isolate test cases stack-info: PR: https://github.com/pytorch/ao/pull/1537, branch: drisspg/stack/28 --- test/test_ops.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 23c0cb938c..a3471d9b5f 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -129,9 +129,16 @@ def test_quant_llm_linear_correctness( TEST_CONFIGS_DEQUANT = list(itertools.product(SHAPES, INNERKTILES, QGROUP_SIZES)) +def make_test_id(param): + if isinstance(param, tuple) and len(param) == 2: # This is a shape + return f"shape_{param[0]}x{param[1]}" + else: # This is inner_k_tiles + return f"tiles_{param}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") # @pytest.mark.skipif(TORCH_VERSION_AT_LEAST_2_5, reason="weight packing is updated in 2.5+") -@pytest.mark.parametrize("shape, inner_k_tiles", TEST_CONFIGS_UNPACK, ids=str) +@pytest.mark.parametrize("shape, inner_k_tiles", TEST_CONFIGS_UNPACK, ids=make_test_id) def test_unpack_tensor_core_tiled_layout_correctness(shape, inner_k_tiles): N, K = shape assert K % (inner_k_tiles * kTileSizeK) == 0 and N % kTileSizeN == 0 @@ -149,7 +156,7 @@ def test_unpack_tensor_core_tiled_layout_correctness(shape, inner_k_tiles): # TODO: Fix "test_aot_dispatch_dynamic" test failure @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") # @pytest.mark.skipif(TORCH_VERSION_AT_LEAST_2_5, reason="weight packing is updated in 2.5+") -@pytest.mark.parametrize("shape, inner_k_tiles", TEST_CONFIGS_UNPACK, ids=str) +@pytest.mark.parametrize("shape, inner_k_tiles", TEST_CONFIGS_UNPACK, ids=make_test_id) def test_unpack_tensor_core_tiled_layout_op(shape, inner_k_tiles): test_utils = [ "test_schema",