Dropout inside scanned function

[Kajiih]

Hello dear Flax community!

I wanted to add some regularization to a rather complex model involving a scanned function. The issue is that I am not sure how to add the static variable training inside the scan. I tried to add it in the carry but the value become traced and causes an error later. Is it ok to use the variable training that is defined out of the scope of the scanned function?

It might not be the same problem but I am getting nans in literally every metrics since I added dropout to the model.

Also, as far as my understanding goes, it is not related to this discussion (even though it helped me figure out a few things when using scan with dropout): #1253.

[cgarciae]

Hey @Kajiih, I believe you need to declare the training argument as a broadcast in scan's in_axes.

[Kajiih]

Thanks for the answer! So it appears to be the same question as this #2127, but I don't understand very well how the broadcasting works here.
I'm using scan on a function I want to loop that is taking a target module, a carry and an input as arguments, so something like

def f(target, carry, inputs):
    ...

but I can add a fourth argument to that function with the argument that is going to be broadcasted, so it becomes

def f(target, carry, inputs_0, inputs_1):
    ...

and then during the definition of the scan, I use in_axes=(0, nn.broadcast) meaning that inputs_0 is scanned over the axis 0 and inputs_1 is broadcasted?

It looks link I was kind of struggling to understand what was the doc refering to in 'in_axes' description.