1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
|
from torch.fx.proxy import Proxy
from ._compatibility import compatibility
@compatibility(is_backward_compatible=False)
def annotate(val, type):
# val could be either a regular value (not tracing)
# or fx.Proxy (tracing)
if isinstance(val, Proxy):
if val.node.type:
raise RuntimeError(f"Tried to annotate a value that already had a type on it!"
f" Existing type is {val.node.type} "
f"and new type is {type}. "
f"This could happen if you tried to annotate a function parameter "
f"value (in which case you should use the type slot "
f"on the function signature) or you called "
f"annotate on the same value twice")
else:
val.node.type = type
return val
else:
return val
|
