Conversation
Some users maintain libraries of code that is largely trace-able but not script-able. However, some functions may need to be `@torch.jit.script`ed if they contain control flow so the tracer will use the compiler version. This however impacts library start up time as in #33418, so this PR adds a workaround in the form of a `@torch.jit._lazy_script_while_tracing` that will only initialize the compiler if the function is called while actually tracing.
facebook-github-bot
added
the
oncall: jit
eellison
left a comment
eellison
left a comment
There was a problem hiding this comment.
LGTM. One of your tests was wrong so if that's not working this may not work right now.
| def fn2(x): | ||
| return untraceable(x) | ||
|
|
||
| with self.capture_stdout(): |
There was a problem hiding this comment.
why capture stdout? you can @suppress_warnings if warnings are being supressed ?
There was a problem hiding this comment.
there are prints in the test
test/test_jit.py
Outdated
| with self.capture_stdout(): | ||
| traced = torch.jit.trace(fn, [torch.ones(2, 2)]) | ||
|
|
||
| FileCheck().check_not("goodbye").check_not("hello").run(traced.graph) |
There was a problem hiding this comment.
shouldn't this be check and not check_not ?
There was a problem hiding this comment.
yeah, this was only passing since the inline mode was wrong
|
Also, not sure |
|
|
| Returns ``True`` in tracing (if a function is called during the tracing of | ||
| code with ``torch.jit.trace``) and ``False`` otherwise. | ||
| """ | ||
| return torch._C._is_tracing |
There was a problem hiding this comment.
Returns the function, and not the bool torch._C._is_tracing(). _script_if_tracing doesn't work on 1.6.0 but is fixed on master.
5 participants
Stacked PRs
@script#34938 - [jit] Remove stray@scriptSome users maintain libraries of code that is largely trace-able but not
script-able. However, some functions may need to be
@torch.jit.scripted ifthey contain control flow so the tracer will use the compiler version.
This however impacts library start up time as in #33418, so this PR adds
a workaround in the form of a
@torch.jit. _script_if_tracingthat will only initialize the compiler if the function is called while
actually tracing.
Differential Revision: D20569778