Adding a new transposed convolution function to lax

[yang-song]

This PR implements lax.gradient_based_conv_transpose. Compared to conv_transpose, it provides support for output_shape and output_padding. It matches the APIs for transposed convolutions derived from the gradient of a forward convolution, which is common in other deep learning frameworks such as TensorFlow, PyTorch, and Keras. This additional function on transposed convolution can make it much easier to reproduce code written in other (and currently more popular) frameworks.

[froystig]

Thank you!

The changes to lax.py include many file-wide formatting adjustments. Could you undo those? We probably don't want to take those at the moment, and it obscures the main change in the diff. You'll also want to squash commits so that there isn't one commit that makes formatting changes followed by another that undoes them.

[froystig]

I'm curious what @hawkinsp and @mattjj think about this change overall, including whether it is best to add it in lax, and for any review comments as well.

[yang-song]

Thank you!

The changes to lax.py include many file-wide formatting adjustments. Could you undo those? We probably don't want to take those at the moment, and it obscures the main change in the diff. You'll also want to squash commits so that there isn't one commit that makes formatting changes followed by another that undoes them.

Just removed formatting changes (done automatically by my IDE).

[schrute99]

Are there any updates on this? It would be awesome to have this function.

[yang-song]

Pending on @froystig and @hawkinsp. I think the requested changes have been made.

[hawkinsp]

I'm not an expert on this, but I'm wondering what the pros and cons are of introducing a new API endpoint vs adding features like output_shape to the existing conv_transpose. What do you think? Is there a reason we need a new function? Is it conceptually different in some important way?

[yang-song]

Because the meaning of padding in conv_transpose is different from that of padding in gradient_based_conv_transpose. If we merge the APIs, we will either break all existing code using conv_tranpose (if adopting the padding in gradient_based_conv_transpose), or otherwise fail to match the APIs of other frameworks.

[levskaya]

Sorry, I wrote the original at a time when no frameworks really existed in JAX. (Nowadays, I'd probably not even add this function to "lax", since it's strictly a specialization of general convolutions, and delegate these matters to NN frameworks.)

Aside from a pending review of correctness, this mainly comes down to a question of organization:

1. Keep our old conv_transpose and delegate these specialized "conv templates" to frameworks.
2. Axe the old conv_transpose and use this new one to match other frameworks... but we probably have users of the old one that prevents that.
3. Fold the alternative behavior definition into the existing one under an optional flag.
4. Retrospectively I wish I had called the existing one "fractionally strided convolutions" since it's more accurate
   maybe that's what we should do? Rename the old one and add the new if it matches what most people expect of "conv transpose"? This comes at the expense of polluting the lax namespace a bit.

[schrute99]

The issue with 1. might be that at least for Flax and Objax, all the convolution modules are just wrappers of the functions in jax.lax. Also if every Jax framework implements its own transposed convolution, they might not be consistent.

[codeboy5]

Hey is this issue being actively worked on ?

[younesbelkada]

Hi all! Is there a plan to merge this PR? It seems that it is the root cause issue of converting some PyTorch models to JAX/FLAX , would be nice if we can merge it ;)

[ericd-1qbit]

The flax docs sent me here and I'd kindly like to add a +1 on hoping this PR will be merged soon :)

[leiteg]

Also came here from the Flax docs. Is there any other way to use transposed convolutions that are compatible with PyTorch's nn.ConvTranspose2d?