mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
538647fe1f
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63888 Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D30523133 Pulled By: jamesr66a fbshipit-source-id: b04cc0d842a74862f42ecba98b757310cd2ec7b0
22 lines
929 B
Python
22 lines
929 B
Python
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
|