From 58721ca9749e5c20881f8248bfa786fc1624def0 Mon Sep 17 00:00:00 2001 From: tankya2 Date: Wed, 27 Aug 2025 11:08:16 +0800 Subject: [PATCH] Add pytest gpu mark --- tests/test_cuquantum_cutensor_backend.py | 4 ++-- tests/test_cuquantum_cutensor_mpi_backend.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_cuquantum_cutensor_backend.py b/tests/test_cuquantum_cutensor_backend.py index 01d8d6e..c492f9f 100644 --- a/tests/test_cuquantum_cutensor_backend.py +++ b/tests/test_cuquantum_cutensor_backend.py @@ -40,7 +40,7 @@ def build_observable_dict(nqubits): return {"terms": terms} - +@pytest.mark.gpu @pytest.mark.parametrize("nqubits", [1, 2, 5, 10]) def test_eval(nqubits: int, dtype="complex128"): """ @@ -81,7 +81,7 @@ def test_eval(nqubits: int, dtype="complex128"): result_sv_cp, result_tn.statevector.flatten() ), "Resulting dense vectors do not match" - +@pytest.mark.gpu @pytest.mark.parametrize("nqubits", [2, 5, 10]) def test_mps(nqubits: int, dtype="complex128"): """Evaluate MPS with cuQuantum. diff --git a/tests/test_cuquantum_cutensor_mpi_backend.py b/tests/test_cuquantum_cutensor_mpi_backend.py index f3dd1c6..4b2deef 100644 --- a/tests/test_cuquantum_cutensor_mpi_backend.py +++ b/tests/test_cuquantum_cutensor_mpi_backend.py @@ -1,4 +1,4 @@ -# mpirun --allow-run-as-root -np 2 python -m pytest --with-mpi test_cuquantum_cutensor_backend_mpi.py +# mpirun --allow-run-as-root -np 2 python -m pytest --with-mpi test_cuquantum_cutensor_mpi_backend.py import math @@ -43,7 +43,7 @@ def build_observable_dict(nqubits): return {"terms": terms} - +@pytest.mark.gpu @pytest.mark.mpi @pytest.mark.parametrize("nqubits", [1, 2, 5, 7, 10]) def test_eval_mpi(nqubits: int, dtype="complex128"): @@ -86,7 +86,7 @@ def test_eval_mpi(nqubits: int, dtype="complex128"): and result_tn_cp.item() == 0 ), f"Rank {backend.rank}: result_tn_cp should be scalar/array with 0, got {result_tn_cp}" - +@pytest.mark.gpu @pytest.mark.mpi @pytest.mark.parametrize("nqubits", [1, 2, 5, 7, 10]) def test_expectation_mpi(nqubits: int, dtype="complex128"): @@ -177,7 +177,7 @@ def test_expectation_mpi(nqubits: int, dtype="complex128"): and result_tn.item() == 0 ), f"Rank {backend.rank}: expected int array [0], got {result_tn}" - +@pytest.mark.gpu @pytest.mark.mpi @pytest.mark.parametrize("nqubits", [1, 2, 5, 7, 10]) def test_eval_nccl(nqubits: int, dtype="complex128"): @@ -217,7 +217,7 @@ def test_eval_nccl(nqubits: int, dtype="complex128"): and result_tn_cp.item() == 0 ), f"Rank {backend.rank}: result_tn_cp should be scalar/array with 0, got {result_tn_cp}" - +@pytest.mark.gpu @pytest.mark.mpi @pytest.mark.parametrize("nqubits", [1, 2, 5, 7, 10]) def test_expectation_NCCL(nqubits: int, dtype="complex128"):