pytorch
1"""
2APIs related to torch.compile which lazily import torch._dynamo to avoid
3circular dependencies.
4"""
5import functools
6
7
8def _disable_dynamo(fn=None, recursive=True):
9"""
10This API should be only used inside torch, external users should still use
11torch._dynamo.disable. The main goal of this API is to avoid circular
12imports issues that is common while using _dynamo.disable inside torch
13itself.
14
15This API avoids it by lazily importing torch._dynamo from the import time to
16the invocation of the decorated function.
17"""
18if fn is not None:
19
20@functools.wraps(fn)
21def inner(*args, **kwargs):
22import torch._dynamo
23
24return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
25
26return inner
27else:
28# decorator usage like @_disable_dynamo(recursive=False). The resulting
29# object expects the original decorated function as the arg.
30return functools.partial(_disable_dynamo, recursive=recursive)
31