Add the general convolution operation to extensions

[DarrenZhang01]

TF XLA version:
https://www.tensorflow.org/xla/operation_semantics?hl=en#conv_convolution;
JAX version: https://jax.readthedocs.io/en/stable/_autosummary/jax.lax.conv_general_dilated.html

[DarrenZhang01]

tests coming soon.

[akshaym]

Hey @DarrenZhang01, Is this ready to review?

[DarrenZhang01]

Hey @DarrenZhang01, Is this ready to review?

Hi Akshay @akshaym ! Actually not yet, there is a bug I need to deal with which I will mention in today's coming meeting.

[akshaym]

Hey @DarrenZhang01, let me know if I can take a look at this. If there are bugs lets just document them for now and add the code anyway? WDYT?

[DarrenZhang01]

Hey @DarrenZhang01, let me know if I can take a look at this. If there are bugs lets just document them for now and add the code anyway? WDYT?

Thanks very much for the support, Akshay! It is ready for being reviewed.

[akshaym]

You had mentioned a bug - is that still around? Or is that fixed?

[DarrenZhang01]

You had mentioned a bug - is that still around? Or is that fixed?

So the original bug is that the transpose convolution was not supported due to the lack of output shape as I mentioned in our last Friday's meeting. Now I added an output shape evaluation helper.

Besides, I used the test cases from JAX general convolution. Those cases are really broad and TF general conv currently do not cover them all, so I skipped some test cases and added a TODO to expand the test cases later on.

The shape evaluation function does not seem so correct, so I remove it and use the output shape of JAX convolution directly as the input shape for TF convolution in the test cases.

[DarrenZhang01]

Thanks for the review, Akshay!

[DarrenZhang01]

Thanks for the review, Akshay!

Now it is ready for review again. @akshaym

[akshaym]

Thanks for the changes @DarrenZhang01!

Just some small remaining nits!

[DarrenZhang01]

Thanks a lot, Akshay! I revised the places that you mentioned.