Issue 1:
It seems that jit decorators are ignored for nested functions. Consider the following example (simplified from an actual program, where inner_fn is slightly more complex):

import jax
import jax.numpy as jnp
from functools import partial

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

@jax.jit
def outer_fn(x):
print("entering outer_fn now")
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()

Expected behavior: As both functions are decorated with jit, inner_fn should be compiled once and the result re-used by the plain for loop inside of outer_fn. Consequently, the output would show only one emission of "entering inner_fn now".

Actual behavior: The jit decorator for inner_fn seems to be ignored: "entering inner_fn now" is printed 10 times and setting env var JAX_LOG_COMPILES=1 only prints one line for Compiling outer_fn. In short, we observe the same behavior as if the jit decorator were absent for inner_fn.

Why a bug?: 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.

Other considerations:

Using fori_loop: Using a fori_loop results in inner_fn 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 inner_fn was small enough so that we actually were seeing performance hits from cuda kernel launches for inner_fn in this case.

No inline=True for jit decorator of inner_fn: In this case the compilation will compile 10 separate sub-jaxpressions for inner_fn that are not inlined but all invoked via xla_call (when inspecting this via jax.make_jaxpr). We see that the jit decorator is therefore not completely without effect, but this is arguably the worst case (separate xla_call overhead as in the fori_loop AND multiple compilation of inner_fn).

Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

If applicable, include full error messages/tracebacks.

Issue 2:
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 jax.experimental.ode.odeint inside of tfp.optimizer.lbfgs_minimize (cf. TFP-on-JAX L-BFGS). This might sound very particular, but I suppose this will come up more often when using JAX in the context of physics models.

However, jit-compilation is very slow, which impedes scaling-up.

Is there a good reason for slow compilation or is there a way to speed this up?

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., len(str(jax.make_jaxpr(...)(...))), as a function of odeint integration steps. The curves have similar shape and become somewhat flat.

!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):
"""
Time derivative of a linear ODE that decays the state.
"""
return -0.1 * x

def loss(x, n_integration_steps):
"""
Minimizing this function essentially just minimizes the norm of x.
"""
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):
"""
Optimizes the initial condition of the ODE system using L-BFGS.
"""
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)