Get index in flax linen scan

[jakubMitura14]

Hello is there a way to access the index of a current iteration inside the scan that would not be traced ?

I know that I can iterate by multiple arrays of the same first dim shape and can put jnp.arange(loop_length) as one thos would give me index but a traced one ; hence I will not be able to use it for example for indexing etc .

Thanks for help !

[jheek]

An untraced iteration count would essentially mean you have different Python code across loop iterations, which is not possible.

You mention you want to index based on the iteration count. This should be possible but you might have run into issues when trying to slice like this:

some_array[i:i + 3]  # error

JAX needs to be able to statically infer the size of your slices and in this case (i + 3) is a tracer and it's not clever enough to realize the size is always 3. Instead you can do the following:

jax.lax.dynamic_slice_in_dim(some_array, i, 3, axis=0)

I think this way you can write your loop even if i is a tracer

[jakubMitura14]

Thanks for your response @jheek ! However, the problem is deeper. For example:

1. I have some einops rearranging that needs to be done and is known at compile time data independent iteration dependent- einops encode rearrangement in strings - I can do a list of strings (as far as I know Jax do not support strings), and access it by getting the index of a pre-computed list (precomputed list saved in setup step) - but I can not access it using an index.

2. I have an image that I resize in each iteration - the size is data-independent and can be precomputed - However, when it is precomputed, I need to access it, and even using a dynamic slice will give traced array back that I can not use in the jax image resize (at least I did not found how to)

3. there are more problems like that - so, for example, dynamic_slice_in_dim works well, but dynamic_slice, when dim is changing, does not ...

Currently, I am using a workaround

def main_fun_op(self,curried,index):
	...

def select_main_fun(self,curried,indexx):
    def a0(self,curried):
        return self.main_fun_op(curried,0)
    def a1(self,curried):
        return self.main_fun_op(curried,1)
    def a2(self,curried):
        return self.main_fun_op(curried,2)
    def a3(self,curried):
        return self.main_fun_op(curried,3)
    def a4(self,curried):
        return self.main_fun_op(curried,4)
    def a5(self,curried):
        return self.main_fun_op(curried,5)
    return nn.switch(indexx, [ a0,a1,a2,a3,a4,a5 ], self, curried)

In that way, I get the index as non-traced - However it is quite clumsy, gives errors when I try to instantiate any flax module inside main_fun_op and is slow to compile.
However, am I right that this workaround, if not mapped, should not pose performance and differentiation problems (as far as I understand documentation, only one branch should be executed)

Basically, What I need is to access not traced data that is known at compilation time, data independent, and different at each iteration

Thanks!

[jheek]

I think you have a few options:

1. unroll (part of) the loop
2. use a switch statement
3. Use jit with a static_argnum for the iteration count and a normal python loop to do the iteration

What is the fastest option depends on hardware and the specific ops and shapes your are using. For GPUs the first and last option will almost certainly be the fastest

[jakubMitura14]

Thanks!

[guillaumepourcel]

Hello,

Will the switch statement execute every branch? I am thinking of a case where the switch (option 2) is used to index a (very long) list of functions.

I know that when used with vmap, switch executes all the branches; however, within a scan, will it do the same thing?

Thanks for the help

I think you have a few options:

1. unroll (part of) the loop
2. use a switch statement
3. Use jit with a static_argnum for the iteration count and a normal python loop to do the iteration

What is the fastest option depends on hardware and the specific ops and shapes your are using. For GPUs the first and last option will almost certainly be the fastest