Tidied; simplified; generalised ConvTranspose implementation.

[patrick-kidger]

@andyehrenberg I think I've got transposed convolutions sorted out.

WDYT? If you're happy with this then I'll merge this in to the other ongoing feature branch.

The main references I used to implement this were Relation 14 of this reference and these animations.

[andyehrenberg]

Definitely happy with this! I really like the clear explanation of output_padding. Do you think it’s worthwhile to include output_shape as an optional argument (that would have to agree with output_padding, if given), as haiku has done? While this could make model building more intuitive, specifying output_padding enables input images of varying sizes in the case of encoder-decoder style architectures, while output_shape could place too much of a restriction. I don’t think it’d be super important/useful to include it. Flax doesn’t seem to include either of these, as noted in google/flax#1872.

[patrick-kidger]

Excellent, glad you like it!

I agree with you that output_shape doesn't sound super important/useful. I think it's nicer to have only one way to do things.

I'll merge this PR in.