inline=True bug in jax.jit?

[ahuang314]

Greetings,

I have here a toy example of using inline=True when jitting a function:

from jax import jit
from functools import partial

@partial(jit, static_argnums=1, inline=True)
def inner_func(x, i):
    return 5 * x, i + 1

@jit
def outer_func(x):
    i = 0
    for _ in range(3):
        x, i = inner_func(x, i)

    return x, i

Now suppose I run the following code:

x = 3
print(outer_func(x))

on jax==0.5.2, this runs perfectly fine and outputs the correct answer. But on jax==0.6.1, this crashes because the compiler tries to mark the argument i in the second call of inner_func as static, but i becomes a tracer after it is returned by the first call of inner_func. I assumed that setting inline=True would effectively copy-paste inner_func into outer_func so that this problem is sidestepped.

Is this a bug in jax==0.6.1, or am I misunderstanding how inline=True is supposed to work? (Although there are obvious workarounds for this problem in this toy example, I would like to avoid them in my actual code if possible)

OS: Windows
Device: CPU

[jakevdp]

Regardless of the JAX version, I think this code is incorrect.

You've marked i as static, but after the first call to inner_func, i will be traced, because the output of a JIT-compiled function can never be static.

The best solution probably is to avoid marking i as static if the function is going to be used in this way.

[ahuang314]

Gotcha. I thought setting inline=True would make it so that nothing is actually returned from inner_func.