Skip to content

[Tests] Migrate tvm.testing.parameters() to pytest.mark.parametrize #19803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tqchen merged 1 commit into from
Jun 17, 2026
Merged

[Tests] Migrate tvm.testing.parameters() to pytest.mark.parametrize #19803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 17 additions & 16 deletions tests/python/codegen/test_target_codegen_cuda_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,23 +142,24 @@ def main(
tvm.testing.assert_allclose(a.numpy().astype("float16"), b.numpy().astype("float16"))


native_dtype, promoted_dtype, numpytype = tvm.testing.parameters(
("float8_e4m3fn", "float32", "float8_e4m3fn"),
("float8_e4m3fn", "float16", "float8_e4m3fn"),
("float8_e4m3fnx2", "float32x2", "float8_e4m3fn"),
("float8_e4m3fnx2", "float16x2", "float8_e4m3fn"),
("float8_e4m3fnx4", "float32x4", "float8_e4m3fn"),
# Supported via half4 vector type extension in codegen
("float8_e4m3fnx4", "float16x4", "float8_e4m3fn"),
("float8_e5m2", "float32", "float8_e5m2"),
("float8_e5m2", "float16", "float8_e5m2"),
("float8_e5m2x2", "float32x2", "float8_e5m2"),
("float8_e5m2x2", "float16x2", "float8_e5m2"),
("float8_e5m2x4", "float32x4", "float8_e5m2"),
("float8_e5m2x4", "float16x4", "float8_e5m2"),
@pytest.mark.parametrize(
"native_dtype,promoted_dtype,numpytype",
[
("float8_e4m3fn", "float32", "float8_e4m3fn"),
("float8_e4m3fn", "float16", "float8_e4m3fn"),
("float8_e4m3fnx2", "float32x2", "float8_e4m3fn"),
("float8_e4m3fnx2", "float16x2", "float8_e4m3fn"),
("float8_e4m3fnx4", "float32x4", "float8_e4m3fn"),
# Supported via half4 vector type extension in codegen
("float8_e4m3fnx4", "float16x4", "float8_e4m3fn"),
("float8_e5m2", "float32", "float8_e5m2"),
("float8_e5m2", "float16", "float8_e5m2"),
("float8_e5m2x2", "float32x2", "float8_e5m2"),
("float8_e5m2x2", "float16x2", "float8_e5m2"),
("float8_e5m2x4", "float32x4", "float8_e5m2"),
("float8_e5m2x4", "float16x4", "float8_e5m2"),
],
)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_fp8_vector_conversions(native_dtype, promoted_dtype, numpytype):
Expand Down
5 changes: 4 additions & 1 deletion tests/python/ir/test_datatype_nv_fp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
# ruff: noqa: F401
import numpy as np
import pytest

import tvm
import tvm.testing
Expand All @@ -29,9 +30,10 @@
float4_e2m1fn = None


np_dtype, dtype_str = tvm.testing.parameters((float4_e2m1fn, "float4_e2m1fn"))
nv_fp4_dtypes = [(float4_e2m1fn, "float4_e2m1fn")]


@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
def test_create_nv_fp4_nd_array(np_dtype, dtype_str):
if np_dtype is None:
"""Skip test if ml_dtypes is not installed"""
Expand All @@ -42,6 +44,7 @@ def test_create_nv_fp4_nd_array(np_dtype, dtype_str):
np.testing.assert_equal(x_nd.numpy(), x)


@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp4_dtypes)
def test_nv_fp4_buffer(np_dtype, dtype_str):
m = te.size_var("m")
n = te.size_var("n")
Expand Down
8 changes: 6 additions & 2 deletions tests/python/ir/test_datatype_nv_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
# ruff: noqa: E501, F401
import numpy as np
import pytest

import tvm
import tvm.testing
Expand Down Expand Up @@ -69,7 +70,7 @@ def func(
return func


np_dtype, dtype_str = tvm.testing.parameters(
nv_fp8_dtypes = [
(float8_e3m4, "float8_e3m4"),
(float8_e4m3, "float8_e4m3"),
(float8_e4m3b11fnuz, "float8_e4m3b11fnuz"),
Expand All @@ -78,9 +79,10 @@ def func(
(float8_e5m2, "float8_e5m2"),
(float8_e5m2fnuz, "float8_e5m2fnuz"),
(float8_e8m0fnu, "float8_e8m0fnu"),
)
]


@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
def test_create_nv_fp8_nd_array(np_dtype, dtype_str):
if np_dtype is None:
"""Skip test if ml_dtypes is not installed"""
Expand All @@ -91,6 +93,7 @@ def test_create_nv_fp8_nd_array(np_dtype, dtype_str):
np.testing.assert_equal(x_nd.numpy(), x)


@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
def test_fp8_unary_op(np_dtype, dtype_str):
func = fp8_unary(dtype_str)
if not tvm.testing.device_enabled("llvm"):
Expand Down Expand Up @@ -119,6 +122,7 @@ def test_fp8_unary_op(np_dtype, dtype_str):
np.testing.assert_equal(args[6].numpy(), expected_a_roundtrip)


@pytest.mark.parametrize("np_dtype,dtype_str", nv_fp8_dtypes)
def test_nv_fp8_buffer(np_dtype, dtype_str):
m = te.size_var("m")
n = te.size_var("n")
Expand Down
24 changes: 20 additions & 4 deletions tests/python/relax/test_op_binary.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: r
tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)


(binary_arith_op, tir_arith_op) = tvm.testing.parameters(
binary_arith_ops = [
(relax.op.add, tirx.Add),
(relax.op.divide, tirx.Div),
(relax.op.floor_divide, tirx.FloorDiv),
Expand All @@ -78,9 +78,10 @@ def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: r
(relax.op.minimum, tirx.Min),
(relax.op.mod, tirx.Mod),
(relax.op.floor_mod, tirx.FloorMod),
)
]


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_arith_infer_struct_info(binary_arith_op: Callable):
bb = relax.BlockBuilder()
vdevice0 = VDevice("llvm")
Expand Down Expand Up @@ -125,6 +126,7 @@ def test_binary_arith_infer_struct_info(binary_arith_op: Callable):
)


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_infer_struct_info_binary_arith_prim_value_with_tensor(binary_arith_op: Callable):
bb = relax.BlockBuilder()

Expand All @@ -134,6 +136,7 @@ def test_infer_struct_info_binary_arith_prim_value_with_tensor(binary_arith_op:
_check_inference(bb, binary_arith_op(x, y), relax.TensorStructInfo((2, 3), "float32"))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_infer_struct_info_binary_arith_prim_value_with_prim_value(binary_arith_op: Callable):
bb = relax.BlockBuilder()

Expand All @@ -143,6 +146,7 @@ def test_infer_struct_info_binary_arith_prim_value_with_prim_value(binary_arith_
_check_inference(bb, binary_arith_op(x, y), relax.PrimStructInfo("float32"))


@pytest.mark.parametrize("binary_arith_op,tir_arith_op", binary_arith_ops)
@pytest.mark.xfail(reason="Not yet implemented")
def test_infer_struct_info_binary_arith_known_prim_value_with_prim_value(
binary_arith_op: Callable, tir_arith_op
Expand All @@ -159,16 +163,17 @@ def test_infer_struct_info_binary_arith_known_prim_value_with_prim_value(
_check_inference(bb, binary_arith_op(y, x), relax.PrimStructInfo(value=tir_y + tir_x))


(binary_cmp_op, tir_cmp_op) = tvm.testing.parameters(
binary_cmp_ops = [
(relax.op.equal, tirx.EQ),
(relax.op.greater, tirx.GT),
(relax.op.greater_equal, tirx.GE),
(relax.op.less, tirx.LT),
(relax.op.less_equal, tirx.LE),
(relax.op.not_equal, tirx.NE),
)
]


@pytest.mark.parametrize("binary_cmp_op", [row[0] for row in binary_cmp_ops])
def test_binary_cmp_infer_struct_info(binary_cmp_op: Callable):
bb = relax.BlockBuilder()
vdev0 = VDevice("llvm")
Expand All @@ -185,6 +190,7 @@ def test_binary_cmp_infer_struct_info(binary_cmp_op: Callable):
_check_inference(bb, binary_cmp_op(x, y2), relax.TensorStructInfo((2, 3), "bool", vdev0))


@pytest.mark.parametrize("binary_cmp_op", [row[0] for row in binary_cmp_ops])
def test_infer_struct_info_binary_cmp_prim_value_to_tensor(binary_cmp_op: Callable):
bb = relax.BlockBuilder()
x = relax.Var("x", R.Tensor((2, 3), "float32"))
Expand All @@ -193,6 +199,7 @@ def test_infer_struct_info_binary_cmp_prim_value_to_tensor(binary_cmp_op: Callab
_check_inference(bb, binary_cmp_op(y, x), relax.TensorStructInfo((2, 3), "bool"))


@pytest.mark.parametrize("binary_cmp_op", [row[0] for row in binary_cmp_ops])
def test_infer_struct_info_binary_cmp_prim_value_to_prim_value(binary_cmp_op: Callable):
bb = relax.BlockBuilder()
x = relax.Var("x", R.Prim("float32"))
Expand All @@ -201,6 +208,7 @@ def test_infer_struct_info_binary_cmp_prim_value_to_prim_value(binary_cmp_op: Ca
_check_inference(bb, binary_cmp_op(y, x), relax.PrimStructInfo("bool"))


@pytest.mark.parametrize("binary_cmp_op,tir_cmp_op", binary_cmp_ops)
@pytest.mark.xfail(reason="Not yet implemented")
def test_infer_struct_info_binary_cmp_known_prim_value_to_prim_value(
binary_cmp_op: Callable, tir_cmp_op
Expand All @@ -217,6 +225,7 @@ def test_infer_struct_info_binary_cmp_known_prim_value_to_prim_value(
_check_inference(bb, binary_cmp_op(y, x), relax.PrimStructInfo(value=tir_cmp_op(tir_y, tir_x)))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_infer_struct_info_shape_symbolic(binary_arith_op: Callable):
bb = relax.BlockBuilder()
m = tirx.Var("m", "int64")
Expand Down Expand Up @@ -245,6 +254,7 @@ def test_binary_infer_struct_info_shape_symbolic(binary_arith_op: Callable):
_check_inference(bb, binary_arith_op(x4, y4), relax.TensorStructInfo(dtype="float32", ndim=-1))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_infer_struct_info_shape_var(binary_arith_op: Callable):
bb = relax.BlockBuilder()
s0 = relax.Var("s0", relax.ShapeStructInfo(ndim=2))
Expand All @@ -266,6 +276,7 @@ def test_binary_infer_struct_info_shape_var(binary_arith_op: Callable):
_check_inference(bb, binary_arith_op(x, y4), relax.TensorStructInfo(dtype="float32"))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_arith_infer_struct_info_more_input_dtype(binary_arith_op: Callable):
bb = relax.BlockBuilder()
x0 = relax.Var("x", R.Tensor((2, 3), "float64"))
Expand All @@ -280,6 +291,7 @@ def test_binary_arith_infer_struct_info_more_input_dtype(binary_arith_op: Callab
_check_inference(bb, binary_arith_op(x2, y2), relax.TensorStructInfo((2, 3), "int64"))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_infer_struct_info_shape_unequal_const_int(binary_arith_op: Callable):
bb = relax.BlockBuilder()
x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
Expand All @@ -288,6 +300,7 @@ def test_binary_infer_struct_info_shape_unequal_const_int(binary_arith_op: Calla
bb.normalize(binary_arith_op(x0, y0))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_arith_infer_struct_info_dtype_mismatch(binary_arith_op: Callable):
bb = relax.BlockBuilder()
x = relax.Var("x", R.Tensor((2, 3), "float32"))
Expand All @@ -296,6 +309,7 @@ def test_binary_arith_infer_struct_info_dtype_mismatch(binary_arith_op: Callable
bb.normalize(binary_arith_op(x, y))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_arith_infer_struct_info_vdevice_mismatch(binary_arith_op: Callable):
bb = relax.BlockBuilder()
x = relax.Var("x", R.Tensor((2, 3), "float32", VDevice("llvm")))
Expand All @@ -304,6 +318,7 @@ def test_binary_arith_infer_struct_info_vdevice_mismatch(binary_arith_op: Callab
bb.normalize(binary_arith_op(x, y))


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_wrong_input_number(binary_arith_op: Callable):
x = relax.Var("x", R.Tensor((2, 3), "float32"))

Expand All @@ -315,6 +330,7 @@ def test_binary_wrong_input_number(binary_arith_op: Callable):
binary_arith_op(x, x, x, x)


@pytest.mark.parametrize("binary_arith_op", [row[0] for row in binary_arith_ops])
def test_binary_infer_struct_info_wrong_input_type(binary_arith_op: Callable):
bb = relax.BlockBuilder()
x0 = relax.Var("x", relax.ShapeStructInfo((2, 3)))
Expand Down
Loading
Loading