diff --git a/tools/cudnn_repro/cudnn_repro/utils.py b/tools/cudnn_repro/cudnn_repro/utils.py index 73bf83148..4ed0f3783 100644 --- a/tools/cudnn_repro/cudnn_repro/utils.py +++ b/tools/cudnn_repro/cudnn_repro/utils.py @@ -164,10 +164,9 @@ def flatten_pass_by_value(value: Any) -> list[int]: if isinstance(value, (int, float)): return [int(value)] if isinstance(value, str): - if value.startswith("0x"): - return [int(value, 16)] + base = 16 if value.startswith("0x") else 10 try: - return [int(value)] + return [int(value, base)] except ValueError: return [] if isinstance(value, list): diff --git a/tools/cudnn_repro/tests/test_cudnn_repro_utils.py b/tools/cudnn_repro/tests/test_cudnn_repro_utils.py new file mode 100644 index 000000000..f4fa5503b --- /dev/null +++ b/tools/cudnn_repro/tests/test_cudnn_repro_utils.py @@ -0,0 +1,31 @@ +from cudnn_repro.utils import flatten_pass_by_value + + +def test_flatten_pass_by_value_valid_hex(): + assert flatten_pass_by_value("0x10") == [16] + + +def test_flatten_pass_by_value_valid_decimal(): + assert flatten_pass_by_value("42") == [42] + + +def test_flatten_pass_by_value_malformed_hex_prefix_only(): + assert flatten_pass_by_value("0x") == [] + + +def test_flatten_pass_by_value_malformed_hex_digits(): + assert flatten_pass_by_value("0xZZ") == [] + + +def test_flatten_pass_by_value_malformed_decimal(): + assert flatten_pass_by_value("abc") == [] + + +def test_flatten_pass_by_value_list_with_malformed_entries(): + assert flatten_pass_by_value(["0x", "0x10", "7", "0xZZ"]) == [16, 7] + + +def test_flatten_pass_by_value_none_and_numbers(): + assert flatten_pass_by_value(None) == [] + assert flatten_pass_by_value(3) == [3] + assert flatten_pass_by_value(2.0) == [2]