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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
|
# Owner(s): ["oncall: jit"]
# flake8: noqa
from dataclasses import dataclass, field, InitVar
from hypothesis import given, settings, strategies as st
from torch.testing._internal.jit_utils import JitTestCase
from typing import List, Optional
import sys
import torch
import unittest
from enum import Enum
# Example jittable dataclass
@dataclass(order=True)
class Point:
x: float
y: float
norm: Optional[torch.Tensor] = None
def __post_init__(self):
self.norm = (torch.tensor(self.x) ** 2 + torch.tensor(self.y) ** 2) ** 0.5
class MixupScheme(Enum):
INPUT = ["input"]
MANIFOLD = [
"input",
"before_fusion_projection",
"after_fusion_projection",
"after_classifier_projection",
]
@dataclass
class MixupParams:
def __init__(self, alpha: float = 0.125, scheme: MixupScheme = MixupScheme.INPUT):
self.alpha = alpha
self.scheme = scheme
class MixupScheme2(Enum):
A = 1
B = 2
@dataclass
class MixupParams2:
def __init__(self, alpha: float = 0.125, scheme: MixupScheme2 = MixupScheme2.A):
self.alpha = alpha
self.scheme = scheme
@dataclass
class MixupParams3:
def __init__(self, alpha: float = 0.125, scheme: MixupScheme2 = MixupScheme2.A):
self.alpha = alpha
self.scheme = scheme
# Make sure the Meta internal tooling doesn't raise an overflow error
NonHugeFloats = st.floats(min_value=-1e4, max_value=1e4, allow_nan=False)
class TestDataclasses(JitTestCase):
@classmethod
def tearDownClass(cls):
torch._C._jit_clear_class_registry()
# We only support InitVar in JIT dataclasses for Python 3.8+ because it would be very hard
# to support without the `type` attribute on InitVar (see comment in _dataclass_impls.py).
@unittest.skipIf(sys.version_info < (3, 8), "InitVar not supported in Python < 3.8")
def test_init_vars(self):
@torch.jit.script
@dataclass(order=True)
class Point2:
x: float
y: float
norm_p: InitVar[int] = 2
norm: Optional[torch.Tensor] = None
def __post_init__(self, norm_p: int):
self.norm = (torch.tensor(self.x) ** norm_p + torch.tensor(self.y) ** norm_p) ** (1 / norm_p)
def fn(x: float, y: float, p: int):
pt = Point2(x, y, p)
return pt.norm
self.checkScript(fn, (1.0, 2.0, 3))
# Sort of tests both __post_init__ and optional fields
@settings(deadline=None)
@given(NonHugeFloats, NonHugeFloats)
def test__post_init__(self, x, y):
P = torch.jit.script(Point)
def fn(x: float, y: float):
pt = P(x, y)
return pt.norm
self.checkScript(fn, [x, y])
@settings(deadline=None)
@given(st.tuples(NonHugeFloats, NonHugeFloats), st.tuples(NonHugeFloats, NonHugeFloats))
def test_comparators(self, pt1, pt2):
x1, y1 = pt1
x2, y2 = pt2
P = torch.jit.script(Point)
def compare(x1: float, y1: float, x2: float, y2: float):
pt1 = P(x1, y1)
pt2 = P(x2, y2)
return (
pt1 == pt2,
# pt1 != pt2, # TODO: Modify interpreter to auto-resolve (a != b) to not (a == b) when there's no __ne__
pt1 < pt2,
pt1 <= pt2,
pt1 > pt2,
pt1 >= pt2,
)
self.checkScript(compare, [x1, y1, x2, y2])
def test_default_factories(self):
@dataclass
class Foo(object):
x: List[int] = field(default_factory=list)
with self.assertRaises(NotImplementedError):
torch.jit.script(Foo)
def fn():
foo = Foo()
return foo.x
torch.jit.script(fn)()
# The user should be able to write their own __eq__ implementation
# without us overriding it.
def test_custom__eq__(self):
@torch.jit.script
@dataclass
class CustomEq:
a: int
b: int
def __eq__(self, other: 'CustomEq') -> bool:
return self.a == other.a # ignore the b field
def fn(a: int, b1: int, b2: int):
pt1 = CustomEq(a, b1)
pt2 = CustomEq(a, b2)
return pt1 == pt2
self.checkScript(fn, [1, 2, 3])
def test_no_source(self):
with self.assertRaises(RuntimeError):
# uses list in Enum is not supported
torch.jit.script(MixupParams)
torch.jit.script(MixupParams2) # don't throw
def test_use_unregistered_dataclass_raises(self):
def f(a: MixupParams3):
return 0
with self.assertRaises(OSError):
torch.jit.script(f)
|
