scale activations in DropPath

[eringrant]

Great project! An additional test + a small fix to align DropPath with standard dropout and torchvision.stochastic_depth.

[paganpasta]

Hey, thanks for the PR.

I was thinking maybe the mode values change to 'batch, rows' will also be beneficial. What do you think?

[eringrant]

For the longer term, I think it makes sense to make equinox.nn.dropout upstream more expressive to handle things like mode = global (as here), or tiled dropout (as handled here, for example). (But of course, this repo is not the place to propose changes to another repo. :)

For the shorter term, I prefer an API that refers to axes by index rather than by their expected semantics. So for example, rather than expecting rows to be meaningful for every input value to eqxvision.layers.DropPath.__call__, I prefer to identify by index the axes over which to broadcast the dropout mask (similarly to the tiled dropout implementation above).

[paganpasta]

That makes sense. It might be beneficial to raise an issue at Equinox for this extension.

For the shorter term, I prefer an API that refers to axes by index rather than by their expected semantics...

Agreed. I'd be happy to accept a PR on this.