How does the JIT get the function types?

username: Enamex
cooked: <p>Does it look at the raw source on disk? How can we give it a <em>custom</em> callable type for a function and force it to recompile a copy of the function body with the new types (but a new name)?</p>

username: SplitInfinity
cooked: <blockquote>
<p>Does it look at the raw source on disk?</p>
</blockquote>
<p>Yes. JIT uses the <code>inspect</code> module to look at the source code for the function, and determines the function schema (types of the inputs and outputs) based on type annotations (either Python3 or mypy style).</p>
<blockquote>
<p>How can we give it a  <em>custom</em>  callable type for a function and force it to recompile a copy of the function body with the new types (but a new name)?</p>
</blockquote>
<p>Could you provide a code sample to clarify this question? What is it that you want to do?</p>

username: Enamex
cooked: <aside class="quote no-group quote-modified" data-username="SplitInfinity" data-post="2" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/5e9695/48.png" class="avatar"> SplitInfinity:</div>
<blockquote>
<blockquote>
<p>How can we give it a <em>custom</em> callable type for a function and force it to recompile a copy of the function body with the new types (but a new name)?</p>
</blockquote>
<p>Could you provide a code sample to clarify this question? What is it that you want to do?</p>
</blockquote>
</aside>
<p>I’d like to be able to say something like:</p>
<pre data-code-wrap="py"><code class="lang-py">def generic_loop(self, inputs, state):
    " ^^^ I'm trying to reuse this. ^^^ "
    outs = []
    for x in inputs:
        out, state = self.fwd(x, state)
        outs.append(out)
    return outs, state
class Cell(nn.Module):
    def __init__(self, forward_sig):
        self.forward = apply_cell_sig(generic_loop, forward_sig)
        # self.forward : (TCustomCell, List[TCustomInput], TCustomState) -&gt; TCustomOut

class CustomCell(Cell):
    def __init__(self):
        super().__init__(
            get_cell_sig(
                CustomCell, TCustomInput, TCustomState, TCustomOut))
    def fwd(self, input: TCustomInput, state: TCustomState):
        return blah(input, state)
</code></pre>
<p>More generally, I’m looking for a way to treat a function body as a “template” that I can provide specific types for, and ask the JIT explicitly to give me a new version for it. So it doesn’t just see its full name and go “I’ve already seen this; here’s the cached version”.</p>
<p>By your confirmation though, this’s probably very difficult. I couldn’t get the JIT to accept a function that takes in a <code>nn.Module</code> instance as an argument either (besides <code>self.</code> accesses) so there’re multiple issues here.</p>

username: SplitInfinity
cooked: <blockquote>
<p>More generally, I’m looking for a way to treat a function body as a “template” that I can provide specific types for, and ask the JIT explicitly to give me a new version for it. So it doesn’t just see its full name and go “I’ve already seen this; here’s the cached version”.</p>
</blockquote>
<p>I don’t think it is a publicly advertised feature, but you could try using <code>@torch.jit._overload</code>. You can find examples of how to use this decorator in <code>test/test_jit.py</code>, like <code>test_function_overloads</code>. The caveat here is that the one body you supply will be used with all overloaded signatures because there is no concept of <code>TypeVar</code> in the JIT typing system. So you might have do type refinement to write one body that works for all types (see <code>my_conv</code> in <code>test_function_overloading_isinstance</code>, also in <code>test_jit.py</code>).</p>
<blockquote>
<p>By your confirmation though, this’s probably very difficult. I couldn’t get the JIT to accept a function that takes in a  <code>nn.Module</code>  instance as an argument either (besides  <code>self.</code>  accesses) so there’re multiple issues here.</p>
</blockquote>
<p>Yeah, modules cannot be passed around because there is no <code>Module</code> type in the JIT type system. There is no way to annotate a function correctly to make this work.</p>

username: Enamex
cooked: <aside class="quote no-group" data-username="SplitInfinity" data-post="4" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/5e9695/48.png" class="avatar"> SplitInfinity:</div>
<blockquote>
<p>I don’t think it is a publicly advertised feature, but you could try using <code>@torch.jit._overload</code> . You can find examples of how to use this decorator in <code>test/test_jit.py</code> , like <code>test_function_overloads</code> . The caveat here is that the one body you supply will be used with all overloaded signatures because there is no concept of <code>TypeVar</code> in the JIT typing system. So you might have do type refinement to write one body that works for all types (see <code>my_conv</code> in <code>test_function_overloading_isinstance</code> , also in <code>test_jit.py</code> ).</p>
</blockquote>
</aside>
<p>This looks promising!</p>
<p>Can you elaborate on what’s happening here?</p>
<pre data-code-wrap="py"><code class="lang-py">        # TODO: pyflakes currently does not compose @overload annotation with other
        # decorators. This is fixed on master but not on version 2.1.1.
        # Next version update remove noqa and add @typing.overload annotation

        @torch.jit._overload  # noqa: F811
        def test_simple(x1):  # noqa: F811
            # type: (int) -&gt; int
            pass

        @torch.jit._overload  # noqa: F811
        def test_simple(x1):  # noqa: F811
            # type: (float) -&gt; float
            pass

        def test_simple(x1):  # noqa: F811
            return x1

        def invoke_function():
            return test_simple(1.0), test_simple(.5)

        self.checkScript(invoke_function, ())

        # testing that the functions are cached
        compiled_fns_1 = torch.jit._script._get_overloads(test_simple)
        compiled_fns_2 = torch.jit._script._get_overloads(test_simple)
        # ^^^ HERE HERE HERE ^^^ #
        for a, b in zip(compiled_fns_1, compiled_fns_2):
            self.assertIs(a.graph, b.graph)
</code></pre>
<p>Why could the two successive invocations of <code>torch.jit._script._get_overloads(test_simple)</code> get different results? At that point, there’s only one thing attached to the <code>test_simple</code> name.</p>

username: Enamex
cooked: <aside class="quote no-group" data-username="SplitInfinity" data-post="4" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/5e9695/48.png" class="avatar"> SplitInfinity:</div>
<blockquote>
<p>Yeah, modules cannot be passed around because there is no <code>Module</code> type in the JIT type system. There is no way to annotate a function correctly to make this work.</p>
</blockquote>
</aside>
<p>I was hoping there’d be a way to work around this specifically <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"> The JIT already works with a <code>self</code> argument of an <code>nn.Module</code> subtype. How can we make this available for free functions? This restriction greatly limits the composability of free functions if I want to use <code>script</code>. Honestly not sure at that point if there’d be a worthwhile speed either. This’s unrelated, but I’ve had some trouble understanding what the JIT does with what it sees, and intuiting where my code might be giving it more trouble than it can handle.</p>

username: tom
cooked: <p>Instead of that, you could also look into</p>
<ul>
<li>Inject your code into linecache under different module names like IPython does. This isn’t exactly documented, but clean as far as the JIT is concerned.</li>
<li>Get the AST using <code>git_jit_def</code>, meddle with the AST, and use <code>_jit_script_compile</code> to compile yourself (you probably need the resolver too), so you basically mimic torch.jit.script except the caching. This is fairly invasive into JIT internals, but at least it is straightforward to your cause.</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

username: tom
cooked: <aside class="quote no-group" data-username="Enamex" data-post="6" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/enamex/48/3329_2.png" class="avatar"> Enamex:</div>
<blockquote>
<p>I was hoping there’d be a way to work around this specifically <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"> The JIT already works with a <code>self</code> argument of an <code>nn.Module</code> subtype. How can we make this available for free functions?  This restriction greatly limits the composability of free functions if I want to use <code>script</code> .</p>
</blockquote>
</aside>
<p>The short answer is not really. The JIT supports classes insofar as they are static - if you JIT-compile Modules, you don’t give it the class source but rather an instance. It will then go through the data members and see and process their types, assuming they will be fixed (which isn’t the case in Python in general).<br>
Maybe you might find more success in taking the “JIT” part more literally, e.g. JIT-compiling a local function from a dynamic function.</p>
<aside class="quote no-group">
<blockquote>
<p>Honestly not sure at that point if there’d be a worthwhile speed either. This’s unrelated, but I’ve had some trouble understanding what the JIT does with what it sees, and intuiting where my code might be giving it more trouble than it can handle.</p>
</blockquote>
</aside>
<p>I’m affraid the JIT isn’t a magic bullet for optimization but “only” does specific things that are highly desirable.</p>

username: Enamex
cooked: <aside class="quote no-group" data-username="tom" data-post="8" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>Maybe you might find more success in taking the “JIT” part more literally, e.g. JIT-compiling a local function from a dynamic function.</p>
</blockquote>
</aside>
<p>Can you elaborate on this a bit more?</p>
<aside class="quote no-group" data-username="tom" data-post="8" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>I’m affraid the JIT isn’t a magic bullet for optimization but “only” does specific things that are highly desirable.</p>
</blockquote>
</aside>
<p>Absolutely. I meant that what those “specific things” are is a bit unclear to me. And doesn’t seem to be a goal of the documentation.</p>

username: SplitInfinity
cooked: <aside class="quote no-group" data-username="Enamex" data-post="5" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/enamex/48/3329_2.png" class="avatar"> Enamex:</div>
<blockquote>
<p>Why could the two successive invocations of <code>torch.jit._script._get_overloads(test_simple)</code> get different results? At that point, there’s only one thing attached to the <code>test_simple</code> name.</p>
</blockquote>
</aside>
<p>There are three - the two overloads and the implementation. And as you can see in the test, they should not have different results in terms of the ordering of functions nor the graphs produced by compiling said functions.</p>

username: tom
cooked: <aside class="quote no-group" data-username="Enamex" data-post="9" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/enamex/48/3329_2.png" class="avatar"> Enamex:</div>
<blockquote>
<p>Can you elaborate on this a bit more?</p>
</blockquote>
</aside>
<p>Well, so one thing could be to actually make it a method of the class of self before you JIT that class / the method.</p>
<aside class="quote no-group" data-username="Enamex" data-post="9" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/enamex/48/3329_2.png" class="avatar"> Enamex:</div>
<blockquote>
<p>Absolutely. I meant that what those “specific things” are is a bit unclear to me. And doesn’t seem to be a goal of the documentation.</p>
</blockquote>
</aside>
<p>My chance to sell advanced PyTorch courses. Your chance to be a hero. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>More seriously we discuss this a bit in section 15.3.1 Interacting with the PyTorch JIT / What to expect from moving beyond classic Python/PyTorch of <a href="https://pytorch.org/deep-learning-with-pytorch" rel="noopener nofollow ugc">out book that you can download in exchange for anwering a few questions</a>.<br>
To summarize that, the main use-cases I see are</p>
<ul>
<li>exporting stuff,</li>
<li>getting rid of the GIL for nicer multithreading in deployment,</li>
<li>optimize certain patterns (e.g. pointwise ops in RNNs and elsewhere are one thing we had relatively early), but people are working on expanding this (e.g. to cover reductions).</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

username: Enamex
cooked: <p>There are three bodies in the text, but I don’t see why the two <code>torch.jit._script._get_overloads(test_sample)</code> calls <em>could</em> return different compiled objects <em>even if</em> the cache was disabled.</p>

username: Enamex
cooked: <aside class="quote no-group" data-username="tom" data-post="8" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>The short answer is not really. The JIT supports classes insofar as they are static - if you JIT-compile Modules, <strong>you don’t give it the class source but rather an instance</strong>. It will then go through the data members and see and process their types, assuming they will be fixed (which isn’t the case in Python in general).</p>
</blockquote>
</aside>
<p>This bit’s interesting. I might’ve been misunderstanding the cache behavior here. If I <code>script</code> different instances of the same class at different times, does it use the cache or treat every instance separately? <code>script</code>, not <code>trace</code>. Because if it treats instances separately, could assigning different modules to attributes be a workaround for JITing functions that make use of modules from arguments (by not passing them through arguments…)?</p>
<aside class="quote no-group" data-username="tom" data-post="8" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>Maybe you might find more success in taking the “JIT” part more literally, e.g. JIT-compiling a local function from a dynamic function.</p>
</blockquote>
</aside>
<p>Which is why I’d imagine being able to “JIT” the same function with different types. Am I misinterpreting this?</p>
<aside class="quote no-group" data-username="tom" data-post="11" data-topic="93112">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>More seriously we discuss this a bit in section 15.3.1 Interacting with the PyTorch JIT / What to expect from moving beyond classic Python/PyTorch of <a href="https://pytorch.org/deep-learning-with-pytorch" rel="noopener nofollow ugc">out book that you can download in exchange for anwering a few questions </a>.</p>
</blockquote>
</aside>
<p>Looks great! Thanks for the rec.</p>

username: SplitInfinity
cooked: <blockquote>
<p>If I  <code>script</code>  different instances of the same class at different times, does it use the cache or treat every instance separately?</p>
</blockquote>
<p>Instances are not cached, but their types in the JIT system are. Every time you script an <code>nn.Module</code>, a type is created for in the JIT type system. This is reused if multiple instances of the same module are scripted in the same program, or even the same instance is scripted twice. The <code>ScriptModule</code> returned by <code>torch.jit.script</code> is always fresh and never from a cache, but it might refer to a type object that is cached.</p>