Add transpose_kernel argument to ConvTranspose.

[mathisgerdes]

This is a proposal for #2577, adding a transpose_kernel option to ConvTranspose such that the previous behavior is maintained by default.

• Add an additional transpose_kernel parameter to ConvTranspose which is passed to jax.lax.conv_transpose
• For padding='CIRCULAR': To maintain the previous convention of ConvTranspose, the alignment of the kernel is changed depending on transpose_kernel.

  # Divide the padding equally between left and right. The choice to put
  # "+1" on the left (or on the right) represents a convention for
  # aligning even-sized kernels.
  if self.transpose_kernel:
    # match convention in Conv to achieve the proper transpose
    total_pad = [
       (size_diff // 2, (size_diff + 1) // 2) for size_diff in size_diffs
    ]
  else:
    # keep the previous convention for backward compatibility
    total_pad = [
       ((size_diff + 1) // 2, size_diff // 2) for size_diff in size_diffs
    ]

[codecov-commenter]

Codecov Report

Merging #2578 (062eed3) into main (7d21975) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #2578      +/-   ##
==========================================
+ Coverage   79.50%   79.52%   +0.01%     
==========================================
  Files          49       49              
  Lines        5206     5211       +5     
==========================================
+ Hits         4139     4144       +5     
  Misses       1067     1067              

Impacted Files	Coverage Δ
flax/linen/linear.py	97.55% <100.00%> (+0.04%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[cgarciae]

Maybe we could move this comment to the else branch and have a similar one for the positive branch?

[mathisgerdes]

I agree, the comments there could be clearer. How about this modified version, with clearer comments for both cases?

[cgarciae]

Thanks for improving this! I have an additional comment, since else this is the default case can we keep the original explanation here? On the positive branch, can you explain why the +1 is now on the right? (e.g. "since the kernel is being transposed we add +1 on the right side instead to mirror the regular version".

[cgarciae]

Thanks @mathisgerdes for doing this! Exposing transpose_kernel is a great idea.

[marcvanzee]

@mathisgerdes @cgarciae this PR has build failures so I have removed the "pull ready" label. Could you please fix those first?

[mathisgerdes]

@mathisgerdes @cgarciae this PR has build failures so I have removed the "pull ready" label. Could you please fix those first?

I believe those were problems with the previous commit, not any changes here. I've rebased to the latest commit.

[mathisgerdes]

Is there still something to be done here, @cgarciae, @andsteing?

[cgarciae]

Sorry for the delay. LGTM.
Thanks for doing this!