1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
# Owner(s): ["module: onnx"]
"""Unit tests for the _core module."""
from __future__ import annotations
import numpy as np
import torch
from torch.onnx._internal.exporter import _core
from torch.testing._internal import common_utils
@common_utils.instantiate_parametrized_tests
class TorchTensorTest(common_utils.TestCase):
@common_utils.parametrize(
"dtype, np_dtype",
[
(torch.bfloat16, np.uint16),
(torch.bool, np.bool_),
(torch.complex128, np.complex128),
(torch.complex64, np.complex64),
(torch.float16, np.float16),
(torch.float32, np.float32),
(torch.float64, np.float64),
(torch.float8_e4m3fn, np.uint8),
(torch.float8_e4m3fnuz, np.uint8),
(torch.float8_e5m2, np.uint8),
(torch.float8_e5m2fnuz, np.uint8),
(torch.int16, np.int16),
(torch.int32, np.int32),
(torch.int64, np.int64),
(torch.int8, np.int8),
(torch.uint16, np.uint16),
(torch.uint32, np.uint32),
(torch.uint64, np.uint64),
(torch.uint8, np.uint8),
],
)
def test_numpy_returns_correct_dtype(self, dtype: torch.dtype, np_dtype):
tensor = _core.TorchTensor(torch.tensor([1], dtype=dtype))
self.assertEqual(tensor.numpy().dtype, np_dtype)
self.assertEqual(tensor.__array__().dtype, np_dtype)
self.assertEqual(np.array(tensor).dtype, np_dtype)
@common_utils.parametrize(
"dtype",
[
(torch.bfloat16),
(torch.bool),
(torch.complex128),
(torch.complex64),
(torch.float16),
(torch.float32),
(torch.float64),
(torch.float8_e4m3fn),
(torch.float8_e4m3fnuz),
(torch.float8_e5m2),
(torch.float8_e5m2fnuz),
(torch.int16),
(torch.int32),
(torch.int64),
(torch.int8),
(torch.uint16),
(torch.uint32),
(torch.uint64),
(torch.uint8),
],
)
def test_tobytes(self, dtype: torch.dtype):
tensor = _core.TorchTensor(torch.tensor([1], dtype=dtype))
self.assertEqual(tensor.tobytes(), tensor.numpy().tobytes())
if __name__ == "__main__":
common_utils.run_tests()
|
