<p dir="auto">It seems that <code class="notranslate">jit</code> decorators are ignored for nested functions. Consider the following example (simplified from an actual program, where <code class="notranslate">inner_fn</code> is slightly more complex):</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, inline=True)
def inner_fn(state):
    print(&quot;entering inner_fn now&quot;)
    return 2*state

@jax.jit
def outer_fn(x):
    print(&quot;entering outer_fn now&quot;)
    old_x = x
    for _ in range(10):
        x = inner_fn(x)
    x = x + old_x
    return x

state = jnp.arange(5, dtype=jnp.uint32)
outer_fn(state).block_until_ready()"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>

<span class="pl-en">@<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>, <span class="pl-s1">inline</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">inner_fn</span>(<span class="pl-s1">state</span>):
    <span class="pl-en">print</span>(<span class="pl-s">"entering inner_fn now"</span>)
    <span class="pl-k">return</span> <span class="pl-c1">2</span><span class="pl-c1">*</span><span class="pl-s1">state</span>

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">outer_fn</span>(<span class="pl-s1">x</span>):
    <span class="pl-en">print</span>(<span class="pl-s">"entering outer_fn now"</span>)
    <span class="pl-s1">old_x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>
    <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">10</span>):
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-en">inner_fn</span>(<span class="pl-s1">x</span>)
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">old_x</span>
    <span class="pl-k">return</span> <span class="pl-s1">x</span>

<span class="pl-s1">state</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">5</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">uint32</span>)
<span class="pl-en">outer_fn</span>(<span class="pl-s1">state</span>).<span class="pl-en">block_until_ready</span>()</pre></div>
<p dir="auto"><strong>Expected behavior</strong>: As both functions are decorated with jit, <code class="notranslate">inner_fn</code> should be compiled once and the result re-used by the plain for loop inside of <code class="notranslate">outer_fn</code>. Consequently, the output would show only one emission of "entering inner_fn now".</p>
<p dir="auto"><strong>Actual behavior</strong>: The jit decorator for <code class="notranslate">inner_fn</code> seems to be ignored: "entering inner_fn now" is printed 10 times and setting env var <code class="notranslate">JAX_LOG_COMPILES=1</code> only prints one line for <code class="notranslate">Compiling outer_fn</code>. In short, we observe the same behavior as if the jit decorator were absent for <code class="notranslate">inner_fn</code>.</p>
<p dir="auto"><strong>Why a bug?</strong>: The current behavior negates any of the caching done for jitted functions, resulting in repeated tracing of a function with identical tracers and thus inflated compilation times and can easily surprise the user.</p>
<p dir="auto"><strong>Other considerations</strong>:</p>
<ol dir="auto">
<li>Using <code class="notranslate">fori_loop</code>: Using a <code class="notranslate">fori_loop</code> results in <code class="notranslate">inner_fn</code> being visited only once, so that will mitigate the issue for this particular example. However, we were actually having troubles with it as it prevents unrolling the loop and our <code class="notranslate">inner_fn</code> was small enough so that we actually were seeing performance hits from cuda kernel launches for <code class="notranslate">inner_fn</code> in this case.</li>
<li>No <code class="notranslate">inline=True</code> for jit decorator of  <code class="notranslate">inner_fn</code>: In this case the compilation will compile 10 separate sub-jaxpressions for <code class="notranslate">inner_fn</code> that are not inlined but all invoked via <code class="notranslate">xla_call</code> (when inspecting this via <code class="notranslate">jax.make_jaxpr</code>). We see that the jit decorator is therefore not completely without effect, but this is arguably the worst case (separate <code class="notranslate">xla_call</code> overhead as in the <code class="notranslate">fori_loop</code> AND  multiple compilation of <code class="notranslate">inner_fn</code>).</li>
</ol>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">I want to optimize the initial condition of a dynamical system w.r.t. a criterion based on the system's trajectory. To this end, I use <code class="notranslate">jax.experimental.ode.odeint</code> inside of <code class="notranslate">tfp.optimizer.lbfgs_minimize</code> (cf. <a href="https://github.com/google/jax/issues/1400#issuecomment-661889919" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/1400/hovercard">TFP-on-JAX L-BFGS</a>). This might sound very particular, but I suppose this will come up more often when using JAX in the context of physics models.</p>
<p dir="auto">However, jit-compilation is very slow, which impedes scaling-up.<br>
Is there a good reason for slow compilation or is there a way to speed this up?</p>
<p dir="auto">Below you see a code example and two graphs that show the compile times and the length of the code generate during jit-compilation, i.e., <code class="notranslate">len(str(jax.make_jaxpr(...)(...)))</code>, as a function of <code class="notranslate">odeint</code> integration steps. The curves have similar shape and become somewhat flat.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="!pip install tfp-nightly==0.12.0.dev20200723

from functools import partial
import jax
import jax.numpy as jnp
from jax.experimental.ode import odeint
from tensorflow_probability.substrates import jax as tfp


def ode_func(x, t):
  &quot;&quot;&quot;
  Time derivative of a linear ODE that decays the state.
  &quot;&quot;&quot;
  return -0.1 * x

def loss(x, n_integration_steps):
  &quot;&quot;&quot;
  Minimizing this function essentially just minimizes the norm of x.
  &quot;&quot;&quot;
  t = jnp.asarray([n*0.1 for n in range(n_integration_steps)])
  y = odeint(ode_func, x, t)
  return jnp.sum(jnp.square(y))  

@partial(jax.jit, static_argnums=(1,))
def optimize(start, n_integration_steps):
  &quot;&quot;&quot;
  Optimizes the initial condition of the ODE system using L-BFGS.
  &quot;&quot;&quot;
  loss_fn = partial(loss, n_integration_steps=n_integration_steps)
  optim_results = tfp.optimizer.lbfgs_minimize(
        jax.value_and_grad(loss_fn),
        initial_position=start,
        num_correction_pairs=10,
        tolerance=1e-8)
  return optim_results.converged, optim_results.position


ndims = 3
n_integration_steps = 10
start = jnp.arange(ndims, 0, -1.)
converged, results = optimize(start, n_integration_steps)"><pre class="notranslate">!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-s1">tfp</span><span class="pl-c1">-</span><span class="pl-s1">nightly</span><span class="pl-c1">==</span><span class="pl-c1">0.12</span>.<span class="pl-c1">0.</span><span class="pl-s1">dev20200723</span>

<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">ode</span> <span class="pl-k">import</span> <span class="pl-s1">odeint</span>
<span class="pl-k">from</span> <span class="pl-s1">tensorflow_probability</span>.<span class="pl-s1">substrates</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span> <span class="pl-k">as</span> <span class="pl-s1">tfp</span>


<span class="pl-k">def</span> <span class="pl-en">ode_func</span>(<span class="pl-s1">x</span>, <span class="pl-s1">t</span>):
  <span class="pl-s">"""</span>
<span class="pl-s">  Time derivative of a linear ODE that decays the state.</span>
<span class="pl-s">  """</span>
  <span class="pl-k">return</span> <span class="pl-c1">-</span><span class="pl-c1">0.1</span> <span class="pl-c1">*</span> <span class="pl-s1">x</span>

<span class="pl-k">def</span> <span class="pl-en">loss</span>(<span class="pl-s1">x</span>, <span class="pl-s1">n_integration_steps</span>):
  <span class="pl-s">"""</span>
<span class="pl-s">  Minimizing this function essentially just minimizes the norm of x.</span>
<span class="pl-s">  """</span>
  <span class="pl-s1">t</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>([<span class="pl-s1">n</span><span class="pl-c1">*</span><span class="pl-c1">0.1</span> <span class="pl-k">for</span> <span class="pl-s1">n</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">n_integration_steps</span>)])
  <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">odeint</span>(<span class="pl-s1">ode_func</span>, <span class="pl-s1">x</span>, <span class="pl-s1">t</span>)
  <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">square</span>(<span class="pl-s1">y</span>))  

<span class="pl-en">@<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>, <span class="pl-s1">static_argnums</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,))</span>
<span class="pl-k">def</span> <span class="pl-en">optimize</span>(<span class="pl-s1">start</span>, <span class="pl-s1">n_integration_steps</span>):
  <span class="pl-s">"""</span>
<span class="pl-s">  Optimizes the initial condition of the ODE system using L-BFGS.</span>
<span class="pl-s">  """</span>
  <span class="pl-s1">loss_fn</span> <span class="pl-c1">=</span> <span class="pl-en">partial</span>(<span class="pl-s1">loss</span>, <span class="pl-s1">n_integration_steps</span><span class="pl-c1">=</span><span class="pl-s1">n_integration_steps</span>)
  <span class="pl-s1">optim_results</span> <span class="pl-c1">=</span> <span class="pl-s1">tfp</span>.<span class="pl-s1">optimizer</span>.<span class="pl-en">lbfgs_minimize</span>(
        <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">loss_fn</span>),
        <span class="pl-s1">initial_position</span><span class="pl-c1">=</span><span class="pl-s1">start</span>,
        <span class="pl-s1">num_correction_pairs</span><span class="pl-c1">=</span><span class="pl-c1">10</span>,
        <span class="pl-s1">tolerance</span><span class="pl-c1">=</span><span class="pl-c1">1e-8</span>)
  <span class="pl-k">return</span> <span class="pl-s1">optim_results</span>.<span class="pl-s1">converged</span>, <span class="pl-s1">optim_results</span>.<span class="pl-s1">position</span>


<span class="pl-s1">ndims</span> <span class="pl-c1">=</span> <span class="pl-c1">3</span>
<span class="pl-s1">n_integration_steps</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
<span class="pl-s1">start</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-s1">ndims</span>, <span class="pl-c1">0</span>, <span class="pl-c1">-</span><span class="pl-c1">1.</span>)
<span class="pl-s1">converged</span>, <span class="pl-s1">results</span> <span class="pl-c1">=</span> <span class="pl-en">optimize</span>(<span class="pl-s1">start</span>, <span class="pl-s1">n_integration_steps</span>)</pre></div>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/26061988/88382886-b058d400-cda9-11ea-8704-0b5580a79005.png"><img src="https://user-images.githubusercontent.com/26061988/88382886-b058d400-cda9-11ea-8704-0b5580a79005.png" alt="compile_times" style="max-width: 100%;"></a> <a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/26061988/88382917-bcdd2c80-cda9-11ea-9c4a-441c917d45e3.png"><img src="https://user-images.githubusercontent.com/26061988/88382917-bcdd2c80-cda9-11ea-9c4a-441c917d45e3.png" alt="jaxpr_str_len" style="max-width: 100%;"></a></p>