test_comparison_utils.py

#!/usr/bin/env python3
# Owner(s): ["module: internals"]


import torch
from torch.testing._internal.common_device_type import (
    instantiate_device_type_tests,
    onlyAccelerator,
)
from torch.testing._internal.common_utils import (
    HardwareClassification,
    run_tests,
    TestCase,
)


class TestComparisonUtils(TestCase):
    hw_classification = HardwareClassification.GENERIC

    def test_all_equal_no_assert(self):
        t = torch.tensor([0.5])
        torch._assert_tensor_metadata(t, [1], [1], torch.float)

    def test_all_equal_no_assert_nones(self):
        t = torch.tensor([0.5])
        torch._assert_tensor_metadata(t, None, None, None)

    def test_assert_dtype(self):
        t = torch.tensor([0.5])

        with self.assertRaises(RuntimeError):
            torch._assert_tensor_metadata(t, None, None, torch.int32)

    def test_assert_strides(self):
        t = torch.tensor([0.5])

        with self.assertRaises(RuntimeError):
            torch._assert_tensor_metadata(t, None, [3], torch.float)

    def test_assert_sizes(self):
        t = torch.tensor([0.5])

        with self.assertRaises(RuntimeError):
            torch._assert_tensor_metadata(t, [3], [1], torch.float)

    def test_assert_layout(self):
        t = torch.tensor([0.5])

        with self.assertRaises(RuntimeError):
            torch._assert_tensor_metadata(t, layout=torch.sparse_coo)


class TestComparisonUtilsDevice(TestCase):
    hw_classification = HardwareClassification.ACCELERATOR

    @onlyAccelerator
    def test_assert_device(self, device):
        t = torch.tensor([0.5], device="cpu")

        with self.assertRaises(RuntimeError):
            torch._assert_tensor_metadata(t, device=device)


instantiate_device_type_tests(
    TestComparisonUtilsDevice, globals(), only_for=("cuda", "xpu"), allow_xpu=True
)

if __name__ == "__main__":
    run_tests()