static_argnums for lax.fori_loop

[malikarjun]

I have the following function

from jax import jit, lax
import jax.numpy as jnp

def multiple_iter():
    def single_iter(i, data):
        data = jnp.pad(data, i)
        return data
    
    data = jnp.zeros((5, 5))
    lax.fori_loop(
        1, 5, single_iter, data
    )
multiple_iter()

This gives ConcretizationTypeError for the line data = jnp.pad(data, i) since i would be an abstract tracer value in single_iter.

I tried to use static_argnums to get around this. The above function is modified as follows

from jax import jit, lax
import jax.numpy as jnp

def multiple_iter():
    def single_iter(i, data):
        data = jnp.pad(data, i)
        return data
    
    data = jnp.zeros((5, 5))
    
    single_iter_jit = jit(single_iter, static_argnums=0)
    lax.fori_loop(
        1, 5, single_iter_jit, data
    )
multiple_iter()

After the modification, I get Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 0) of type <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'> for function single_iter is non-hashable. error. I have the following questions

1. Is it impossible to use the pad function with dynamic arguments?
2. In this case, since single_iter will only be called with a predefined set of values 1, 2, 3, 4, will it still lead to cache misses (assuming Non-hashable static arguments were supported)?

[YouJiacheng]

fori_loop will trace the loop body with tracer, thus single_iter will be called with all tracer argument.
Moreover, fori_loop cannot change the shape of val.
You should use python for loop in such case.