How to scan over list input?

[py4]

I have the following loop implementation for illustration:

class Model(nn.Module):
   num_layers: int
   @nn.compact
   def __call__(self, x: jax.Array, ys: List[int]):
       for i in range(self.num_layers):
            x = Layer()(x, ys[i])

       return x

I like to do this with nn.scan. But I get the impression that it can only scan over variables not additional inputs. I tried this:

Scanner = nn.scan(Layer, split_rngs={'params': True}, in_axes=(0,), out_axes=0, variable_axes={'params': 0}, length=self.num_layers)()
x, _ = Scanner(x, ys)

But it seems like entire ys is being passed to each iteration of loop. Basically the code is not scanning over ys. Any pointer is appreciated!

[cgarciae]

Hey, you need to convert the list into an Array first using jnp.stack, else scan will interpret it as a pytree.