Scan over a module with a different number of inputs and outputs

[hr0nix]

Suppose I have a module that takes some number of inputs, but returns just one output. Something like

class MyModule(nn.Module):
  @nn.compact
  def __call__(self, x, y):
    return x + y

Now suppose I want to use nn.scan to build a stack of such modules, but in such a way that x is carried over the loop, but the value of y is shared between the modules in the stack. So, essentially, my stack should do the following:

for i in range(num_layers):
  x = MyModule()(x, y)

Unfortunately, I cannot figure out from documentation how to achieve that with nn.scan. Any suggestions would be much appreciated.

[cgarciae]

Hey @hr0nix, you can set in_axes to nn.broadcast to indicate that y should be the same on all steps:

class MyModule(nn.Module):
  @nn.compact
  def __call__(self, x, y):
    return x + y, None
    
StackedMyModule = nn.scan(MyModule, in_axes=nn.broadcast)

Note that you must return (carry, output), here I am returning None as the scan output.