Does jax.jit(..., static_argnums=...) work for non-top-level functions?

[dfdx]

Consider a top-level function wrapper calling a lower-level function jnp.triu:

import jax
import jax.numpy as jnp


def wrapper(x: jax.Array, start: int):
    return jnp.triu(x, k=start + 1)


x = jax.random.normal(jax.random.PRNGKey(0), (3, 3))
wrapper(x, 1)   # OK

jnp.triu is decorated with @partial(jit, static_argnames=('k',)), so its second argument is static anyway.

Now I want to JIT-compile wrapper too. However, calling jax.jit(wrapper)(x, 1) fails with the following error:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.
The error occurred while tracing the function triu at /root/.pyenv/versions/3.10.6/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:2850 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:i32[] = sub b c
    from line <ipython-input-43-329900c1d420>:25 (wrapper)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

As far I understand, JIT-compiling the wrapper cancels JIT-compilation of the inner jnp.triu, including effect of static_argnames. In this particular example, I can add a similar argument to compilation of the wrapper itself:

jit_wrapper = jax.jit(wrapper, static_argnums=1)
jit_wrapper(x, 1)    # OK

But in more realistic examples the call stack is much deeper and specifying static_argnames for inner functions on the top level is not feasible.

Is my understanding correct? If so, are there any workarounds to specify static_argnames more "locally"?

[jakevdp]

The issue is: you cannot pass a dynamic argument to a function which expects a static argument.

In your first example, jit(wrapper)(x, 1) results in start being treated as a dynamic variable, and so it cannot be passed to a function that requires a static input. There is no way for the inner jit to communicate to the outer jit which arguments should be static and which should not, and in my opinion it would not be desirable for the API to support that kind of "action at a distance". As they say, explicit is better than implicit, and so you must explicitly specify which arguments are static for each function you wrap in jit.

Does that make sense?

[dfdx]

Correct me if I'm wrong, but static arguments actually can be passed to a JIT-wrapped function, yet they will lead to re-compilation with the new value, right? For example, if I invoke jnp.triu(x, 1) and then jnp.triu(x, 2), it will produce 2 entries in the JIT cache:

In [3]: x = jax.random.normal(jax.random.PRNGKey(0), (3, 3))
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

In [4]: jnp.triu._cache_size()
Out[4]: 0

In [5]: jnp.triu(x, 1)
Out[5]: 
Array([[ 0.        ,  0.26423115, -0.18252768],
       [ 0.        ,  0.        , -0.1521442 ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float32)

In [6]: jnp.triu._cache_size()
Out[6]: 1

In [7]: jnp.triu(x, 2)
Out[7]: 
Array([[ 0.        ,  0.        , -0.18252768],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float32)

In [8]: jnp.triu._cache_size()
Out[8]: 2

In [9]: x2 = jax.random.normal(jax.random.PRNGKey(0), (3, 3))

In [10]: jnp.triu(x2, 2)    # x2 is not static and has the same size as x, so it doesn't trigger re-compilation
Out[10]: 
Array([[ 0.        ,  0.        , -0.18252768],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float32)

In [11]: jnp.triu._cache_size()
Out[11]: 2

Then the question is whether a JIT-compiled function can access this cache & compile missing sets of arguments. From your answer I understand it cannot. Probably, because it would require XLA to call back to the compiler, which is impossible. Is this analysis correct?

If so, could you please clarify one more detail? If I lift static arguments to the wrapper (i.e. jax.jit(wrapper, static_argnums=1)), call it with, say, 1000 different values and get 1000 entries in the compilation cache, should I expect any bad effects beyond increased cache size?

[jakevdp]

Cacheing is only relevant for the outermost JIT, and each unique value of a static argument in the outermost JIT will lead to a new cache entry. Does that answer your question?