ResBlock combinator in stax

[aeftimia]

I have seen (and used) the ResNet-style "split, apply, combine" design pattern enough that I thought it might be nice to have it built in. This may be trivial, but I thought it made code a little cleaner and could potentially clean up import statements for architectures with residual connections.

I confirmed the resnet50 example exhibited identical losses for each update with and without the new ResBlock call.

Please let me know if you like the idea, but would prefer some other function name and/or signature--I would be happy to modify it to suite maintainer preferences.

[shoyer]

Looks reasonable to me, though my preference would be to avoid the final “tail” layer defaulting to ReLU and keep all activations explicit. This would be more consistent with the rest of stax, and in my experience unexpected activations are a common source of bugs.

[shoyer]

You can’t use this syntax yet because JAX still supports Python 2 (for a little while longer)

[aeftimia]

Looks reasonable to me, though my preference would be to avoid the final “tail” layer defaulting to ReLU and keep all activations explicit. This would be more consistent with the rest of stax, and in my experience unexpected activations are a common source of bugs.

Oh interesting--in hindsight I could see that potentially being an issue but I don't think I would have thought of it before you mentioned it. I modified the default tail to be Identity.

You can’t use this syntax yet because JAX still supports Python 2 (for a little while longer)

I submitted a patch that works for Python 2 and 3 with a TODO to replace it with the more readable function head when Python 2 support is dropped.

Please feel free to let me know if you would like any further modifications.

[jekbradbury]

Thanks for the PR! We're leaning towards declining this change even thought it's correct and increases the readability of the ResNet definition, because stax.py is meant more as an example (forking encouraged!) than as a growing or comprehensive library and more full-featured neural network creation APIs are available or coming soon (see https://github.com/google/trax and https://github.com/JuliusKunze/jaxnet for two of them).

[aeftimia]

Thanks for the PR! We're leaning towards declining this change even thought it's correct and increases the readability of the ResNet definition, because stax.py is meant more as an example (forking encouraged!) than as a growing or comprehensive library and more full-featured neural network creation APIs are available or coming soon (see https://github.com/google/trax and https://github.com/JuliusKunze/jaxnet for two of them).

Thanks for the update. I was not aware of either of those two projects. It looks like Trax already has a similar residual layer.

[hawkinsp]

We're no longer evolving stax in favor of other libraries such as Flax, Haiku, and Trax. Closing!