raise ConcretizationTypeError when use fori_loop

[yiminghwang]

Hi, there

I am going to use fori_loop, but it raises the ConcretizationTypeError like,

jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>
The problem arose with the `bool` function.

the simplified code is shown as follows,

import jax

def random_binary_array(key, n):
    key1, key2 = jax.random.split(key, num=2)
    return key1, jax.random.choice(key2, jnp.asarray([0, 1]), shape=(n,))

def loop(i, s):
    key1, key2 = jax.random.split(s["key"], num=2)
    s["key"], bin_array =  random_binary_array(key2,s["n"])
    return s

n = 5
key = jax.random.PRNGKey(0)
key, bin_array = random_binary_array(key, 5)
s = {
    "n": n,
    "key": key,
    "bin_array": bin_array
}

new_bin_array = jax.lax.fori_loop(0,5,loop,s)

does anyone know what is the reason for this error and how to fix it?

[jakevdp]

Within the loop function, all arguments will be traced. This means that s['n'] will be traced, and so when you write shape=(n,) you are requesting a dynamicly-shaped array, which is not currently supported in JAX.

To fix this, you'll need to use a static value for n. One way to do this would be to use the global n value; that is change this line:

s["key"], bin_array = random_binary_array(key2, s['n'])

to this:

s["key"], bin_array = random_binary_array(key2, n)

Of course this assumes that n is not changing from iteration to iteration; if it is, then you really do have dynamic array sizes and you'll have to find a different approach to your algorithm within JAX.

[yiminghwang]

Many thanks for the prompt reply! Oh, I see.
Go a step further, what if I want to save the result in each iteration, instead of creating the global variable to store, is there any other method to achieve it?

[mattjj]

Would jax.lax.scan do what you want?

[yiminghwang]

I see, thanks! I will check it.