From 4df65cb382617972ca687a96a5e5ba926a109ef5 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Mon, 11 Aug 2025 01:07:21 +0800 Subject: [PATCH] Add assertions in tests for conversion. --- tests/conversion_test.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/conversion_test.py b/tests/conversion_test.py index 3454319..6c49f82 100644 --- a/tests/conversion_test.py +++ b/tests/conversion_test.py @@ -6,7 +6,7 @@ @pytest.fixture() def x() -> GrassmannTensor: - return GrassmannTensor((False, False), ((2, 2), (1, 3)), torch.randn([4, 4])) + return GrassmannTensor((False, False), ((2, 2), (1, 3)), torch.randn([4, 4], device="cpu:0")) @pytest.mark.parametrize("dtype_arg", ["position", "keyword", "none"]) @@ -21,7 +21,8 @@ def test_conversion( args: list[typing.Any] = [] kwargs: dict[str, typing.Any] = {} - device = torch.device("cpu") if device_format == "object" else "cpu" + device_str = "cuda:0" if torch.cuda.is_available() else "cpu:0" + device = torch.device(device_str) if device_format == "object" else device_str match device_arg: case "position": args.append(device) @@ -38,8 +39,16 @@ def test_conversion( case _: pass - if len(args) <= 1: - y = x.to(*args, **kwargs) + if len(args) > 1: + pytest.skip("Cannot pass both dtype and device as positional arguments") + + y = x.to(*args, **kwargs) + assert isinstance(y, GrassmannTensor) + assert y.arrow == x.arrow + assert y.edges == x.edges + assert y.tensor.dtype == torch.complex128 if dtype_arg != "none" else torch.float32 + assert y.tensor.device.type == (torch.device(device_str) if device_arg != "none" else torch.device("cpu:0")).type + assert torch.allclose(y.tensor, x.tensor.to(dtype=y.tensor.dtype, device=y.tensor.device)) def test_conversion_duplicated_value(x: GrassmannTensor) -> None: