Allow passing kwargs to the layers in the Sequential layer

[Artur-Galstyan]

Hi there,

I had just noticed that I cannot pass any parameters to the layers defined in the Sequential layer - except for x, state, and key. If the layers within Sequential require additional parameters, then you're out of luck, because x is defined as Array (as opposed to tuple[Any] (e.g.))

I created a simple PR for this, see #652

[patrick-kidger]

Thanks!
What's the use-case? Right now I'm tempted to suggest that if you need this, then to just write things out directly. Sequential is really only designed for simple models.

[Artur-Galstyan]

Concretely, I need this for my Transformer. See line 191

Each Block call takes a mask as input (because it has a MHA). The Transformer has a Sequential with n Block elements. Each Block should get the same mask. Therefore, I want to pass those to each Block.

I don't think we lose anything here and don't make it more complicated. It's up to the user - if there are no kwargs, then simply don't use them, but the option is there if you need to 🙂

I should also mention that the solution is not to change the signature to tuple[any] but rather to allow for kwargs and simply forward those kwargs to each module within Sequential.

[patrick-kidger]

Hmm, I agree it's not too much more. But also just writing out the for loop manually isn't too tricky either. (In fact if I was to go back and do Equinox again, I'd probably leave out Sequential altogether!)

So on that basis I think I'm probably going to suggest keeping things simple and not implementing this, I'm afraid! :)

[Artur-Galstyan]

Ok 🥲 I'll close the PR and this issue then.