pytorch

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

4
import torch
5
from torch.testing._internal.common_utils import TestCase, run_tests
6

7
class TestComparisonUtils(TestCase):
8
    def test_all_equal_no_assert(self):
9
        t = torch.tensor([0.5])
10
        torch._assert_tensor_metadata(t, [1], [1], torch.float)
11

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

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

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

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

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

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

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

34

35
if __name__ == '__main__':
36
    run_tests()
37