Padding in conv_transpose()

[EngineerKhan]

I have made a simple example:

a = jnp.ones((28,28))
b = jnp.ones((3,3))
c = lax.conv_general_dilated(a[None, None], b[None, None], (1,1),[(0,0),(0,0)],(1,1))

But when I try to switch conv_general_dilated() to conv_transpose() (with the same parameters), it gives me error:

File "/usr/local/lib/python3.8/dist-packages/jax/_src/lax/lax.py", line 1861, in conv_transpose
if padding in {'SAME', 'VALID'}:
TypeError: unhashable type: 'list'

While comparing the arguments of both functions (and implementation details as well) respectively, I don't find any difference between the padding handling of the two. How can I fix it?

[jakevdp]

This is a bug that comes from this line: https://github.com/google/jax/blob/9acb7891acda6a8ec12ac0b27d5cf8538eedd958/jax/_src/lax/convolution.py#L322

A set membership check fails if the query is not hashable, and lists are not hashable. We should replace the set {'SAME', 'VALID'} with the list ['SAME', 'VALID'] to address this.

In the meantime, you can work around this by passing your paddings as a tuple rather than a list:

c = lax.conv_general_dilated(a[None, None], b[None, None], (1,1),((0,0),(0,0)),(1,1))

[jakevdp]

#8696 has the fix - thanks!