The Jax primitive jax.lax.conv_transpose result does not match the result of torch.nn.functional.conv_transpose2d #32570
Replies: 4 comments 7 replies
-
|
I fed your question to Gemini 2.5 Pro, and it says (and for the record, I agree with it): Of course, I can help you solve this JAX issue. It's a very common point of confusion when working with transposed convolutions. The core of the problem lies in a misunderstanding of how jax.lax.conv_transpose interprets the dimensions of the kernel (your weight matrix W). Let's break it down. In a Nutshell: The SolutionYou should not transpose your weight matrix W. The error message you are seeing is correct and points to a channel mismatch created by your transpose operation. This is the correct implementation: Did that answer your question? |
Beta Was this translation helpful? Give feedback.
-
|
gentle advice @SanjithKumar2, in your place I wouldn't leave full path on debugging message I share (even if it is generic), just the jax directory bit |
Beta Was this translation helpful? Give feedback.
-
|
Thank you @DiagRisker next time will make sure that! |
Beta Was this translation helpful? Give feedback.
-
|
@hawkinsp Thank you for taking the time to answer my question I tried this already as I mentioned, it did in-fact return the output of right shape (1, 1, 28, 28), but it did not match with the torch result, now to clarify the naming conventions once more Convolution forward pass:X-> (1, 1, 28, 28) (this is the actual input shape during the forward pass) (batch_size, in_channel, H, W) In this case the JAX and Torch results matched perfectly. import torch.nn.functional as F
import jax.lax as L
# Torch forward
Z = F.conv2d(X, W, stride=(2,2)) + 1e-12
# Jax forward
Z_ = L.conv_general_dilated(
lhs=X,
rhs=jnp.transpose(W, (2,3,1,0)),
window_strides = (2, 2),
padding='VALID',
dimension_numbers=('NCHW','HWIO','NCHW')
)
# Comparison
print(torch.allclose( torch.tensor(Z_), Z) # -> prints Truealso compared the forward pass with my manual implementation, both jax and torch matched Now during the backward relevance propagation (CONV TRANSPOSE)Y -> (1, 15, 13, 13) (this is the output relevance score and matches the output shape from the forward pass) (batch_size, out_channel, oH, oK) K, K_ -> (1, 1, 28, 28) (the expected output from conv transpose) (batch_size, in_channel (according to forward pass), H, W) To be clear this is how I see the dimensions and you can correct me here, In this case given your solution (from gemini 2.5 pro) did not match the torch results (the torch results matched my manually calculated results... that is why i said pytorch was correct) # Torch
K = F.conv_transpose2d(Y, W, stride=(2,2)) # output shape -> (1, 1, 28, 28)
# JAX
K_ = L.conv_transpose(
lhs = Y,
rhs = W,
strides = (2,2),
padding='VALID',
dimension_numbers=('NCHW','IOHW','NCHW')
) # output shape -> (1, 1, 28, 28)
# Comparison
print(torch.allclose(K, torch.tensor(K_)) # prints FalseWhen I check it against my manual implementation the pytorch results match, but not the JAX results I got traumatized by ChatGPT and Grok, and I do not have Pro versions for any llm, that is why I came here. Hope this helps you to understand more... |
Beta Was this translation helpful? Give feedback.
-
|
Please provide a minimal, self-contained, runnable example of the problem. i.e, something I can run. |
Beta Was this translation helpful? Give feedback.
-
|
@hawkinsp I want the final print statement to print True import jax.numpy as jnp
import jax.lax as L
import torch
import torch.nn.functional as F
key = jax.random.PRNGKey(0)
# JAX
Y = jax.random.normal(key,(1,15,13,13))
W = jax.random.normal(key,(15, 1, 4, 4))
stride = (2,2)
padding = 'VALID'
K_ = L.conv_transpose(
lhs = Y,
rhs = W,
strides = stride,
padding=padding,
dimension_numbers=('NCHW','IOHW','NCHW')
)
# Torch
Y = torch.tensor(Y)
W = torch.tensor(W)
K = F.conv_transpose2d(input=Y,weight=W,stride=(2,2))
# Compare
K_ = torch.tensor(K_)
print(torch.allclose(K, K_)) |
Beta Was this translation helpful? Give feedback.
-
|
Hello @SanjithKumar2 for dimension_numbers on jax.lax.conv_general_dilated:
your W.shape is incorrect! to change the convention used by jax.lax.conv_general_dilated
rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))
out_spec = lhs_specI remember having troubles setting it right, but looking at flax code: https://github.com/google/flax/blob/main/flax/nnx/nn/linear.py
Here is a simple example of dilation import jax, jax.numpy as jnp, jax.lax as lax
from jax.tree_util import Partial, tree_map
def Varmat(var, L):
""" Variable wrapper """
if not hasattr(var,"__iter__"):
var = (var,)*L
elif len(var)-L:
var = (var[0],)*L# careful to how you declare var
return var
# Here is an example of how you dilate
def dilate( X, dilC = 2):
dilC = Varmat(dilC, X.ndim)
W = jnp.array([1.]).reshape( [1]*X.ndim)
pad = [(0,0)]*X.ndim # VALID padding
stride = (1,)*X.ndim
X = X.reshape((1,1)+X.shape).astype(jnp.float32)
W = W.reshape((1,1)+W.shape)
return lax.conv_general_dilated( X, W, padding = pad, lhs_dilation = dilC, window_strides = stride)[0,0]
X = jnp.ones((11,11))
dil = (3,2)
Xnew = dilate(X, dilC = dil)
from matplotlib import pyplot as plt ; plt.ion()
plt.figure()
plt.imshow(Xnew)
you can also force code the convention of https://docs.pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html with the kernel_dil , input_dil and dimension_numbers parameters In your example the expected (1, 1, 28, 28) is obtained by chosing a dilation factor of 2 on W,H part of the shape, instead of strides. In_f = 15 # input features
Out_f = 1
batchsdim =1
input = jnp.ones( (batchsdim, In_f , 13, 13) ) # n + 2 dimensions (1, 15, 13, 13)
W = jnp.ones( ( Out_f , In_f , 4, 4) ) # n + 2 dimensions ( 1, 15, 4, 4)
dil = (2,2)
def conv_dilate( X, W, dilC = 2):
dilC = Varmat(dilC, 2) # n length
pad = tuple(map( lambda dil: (dil+1,dil+1), dilC )) # SAME padding with dilation # n length
stride = (1,)*2 # n length
return lax.conv_general_dilated( X.astype(jnp.float32) , W.astype(jnp.float32) , padding = pad, lhs_dilation = dilC, window_strides = stride)
out = conv_dilate(input ,W, dilC= dil)
out.shape # (1, 1, 28, 28)This is not a bug, it's just less automated :) |
Beta Was this translation helpful? Give feedback.
-
|
@DiagRisker Thank you for your reply, I tried running your solution and my problem is the final result (out) still does not match the torch result... what I want is whether the result after the calculation is same for both frameworks when i do the same operation with the same set of inputs. # JAX following your implementation
Array([[[[ 9.68568087e-01, 4.70307446e+00, 6.82874203e+00,
3.16216063e+00, 3.14307642e+00, 5.14505196e+00,
3.81126118e+00, 4.61444569e+00, -3.79500413e+00,
4.37498808e-01, 2.16927624e+00, -7.55848408e+00,
1.52375758e+00, -1.04054127e+01, 2.29220128e+00,
5.70161152e+00, 5.07335663e-02, 6.20733261e+00,
-6.10543299e+00, 1.11717920e+01, 9.62757778e+00,
2.75180817e+00, -7.27785587e+00, 3.55764771e+00,
-9.67870474e-01, -8.50762546e-01, -4.12547678e-01,
-2.06366229e+00],
....
-6.27017498e+00, 5.63667583e+00, -1.20715199e+01,
-6.41208363e+00, 7.43806958e-01, -4.47466993e+00,
-8.35354805e+00, -2.25002718e+00, 2.99455690e+00,
-4.33305025e+00, -9.20791149e+00, -9.37315941e-01,
-3.90706158e+00]]]], dtype=float32)
# Torch
tensor([[[[-1.2977e+00, 7.6736e+00, 3.9104e+00, -2.2823e+00, -9.1921e-01,
2.8535e+00, 3.4661e+00, -6.8813e-01, -4.5785e+00, -9.3846e+00,
-3.6988e+00, -5.7721e+00, 1.6770e+00, -9.2062e-01, 3.0225e+00,
4.1085e+00, 5.2094e-01, 7.1012e+00, -7.3374e+00, 5.9097e+00,
-4.3490e+00, 4.5977e+00, -2.8287e-02, 2.1917e+00, 5.9463e+00,
4.1561e+00, -3.4796e+00, -9.1264e-01],
....
1.3091e+00, 6.5275e-01, 1.7290e+00, 2.5780e+00, 2.8498e+00,
-9.3760e-01, 3.1753e+00, 1.1495e+00, 5.3902e-01, 3.2018e-01,
-3.3610e+00, 1.3772e+00, -8.5320e+00, -1.2680e+00, -7.9145e+00,
-9.8107e+00, 1.6005e+00, 5.6556e+00, 4.3627e+00, -5.1006e+00,
1.2330e-01, 6.5773e+00, -7.1382e+00]]]]) |
Beta Was this translation helpful? Give feedback.
-
|
@SanjithKumar2 can you show me how you produced the JAX Array ? (did you keep a stride of 1? assuming the torch.tensor code part didn't change) especially these part if isinstance(padding, str) and padding in {'SAME', 'VALID'}:
if rhs_dilation is None:
rhs_dilation = (1,) * (rhs.ndim - 2)
effective_k_size = map(lambda k, r: core.dilate_dim(k, r), k_sdims, rhs_dilation)
pads = [_conv_transpose_padding(k, s, padding)
for k,s in zip(effective_k_size, strides)]
[...]
if transpose_kernel:
# flip spatial dims and swap input / output channel axes
rhs = _flip_axes(rhs, np.array(dn.rhs_spec)[2:])
rhs = rhs.swapaxes(dn.rhs_spec[0], dn.rhs_spec[1])
# knowing
def _conv_transpose_padding(k, s, padding):
"""Calculate before and after padding for a dim of transposed convolution.
Args:
k: int: kernel dimension.
s: int: dimension stride value.
padding: 'same' or 'valid' padding mode for original forward conv.
Returns:
2-tuple: ints: before and after padding for transposed convolution.
"""
if padding == 'SAME':
pad_len = k + s - 2
if s > k - 1:
pad_a = k - 1
else:
pad_a = int(np.ceil(pad_len / 2))
elif padding == 'VALID':
pad_len = k + s - 2 + max(k - s, 0)
pad_a = k - 1
else:
raise ValueError('Padding mode must be `SAME` or `VALID`.')
pad_b = pad_len - pad_a
return pad_a, pad_bIt shouldn't be very hard to match the code for you needs given https://docs.pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html definition (sadly I have not the time right now) |
Beta Was this translation helpful? Give feedback.
-
|
@DiagRisker The previous Jax result you can ignore as it was for the dilated convolution you gave, but I would suggest you run the snippet I shared earlier... when i run it the results do not match |
Beta Was this translation helpful? Give feedback.
-
|
@SanjithKumar2 in JAX conv_transpose macro we have definition swap (according to deconvolution's one): conv_general_dilated(lhs, rhs, window_strides = one, padding = pads, lhs_dilation = strides, rhs_dilation = rhs_dilation,
dimension_numbers = dn, precision=precision, preferred_element_type=preferred_element_type)As shown upper, you only use stride (ergo lhs_dilation in practice here), so we just have to transpose the "spatial dimensions" import jax, jax.numpy as jnp, jax.lax as lax
from jax.tree_util import Partial, tree_map
from numpy.random import randint
def Varmat(var, L):
""" Variable wrapper """
if not hasattr(var,"__iter__"):
var = (var,)*L
elif len(var) - L:
var = (var[0],) * L
return var
class Rng(): # for manual initialization
def __init__(self, seed = randint(2**30)):
print('seed :', seed,';')
self.key = jax.random.key(seed)
def __call__(self, n = ()):
R = jax.random.split(self.key, n)
self.key = R.flatten()[0]
return R
rky = Rng()
In_f = 15 # input features
Out_f = 1
batchsdim =1
Y= jax.random.normal( rky(), shape = (batchsdim, In_f , 13, 13) ) # n + 2 dimensions (1, 15, 13, 13)
W_ = jax.random.normal(rky(), shape = ( Out_f , In_f , 4, 4) ) # n + 2 dimensions ( 1, 15, 4, 4)
dil = (2,2)
def conv_dilate( X, W, dilC = 2):
dilC = Varmat(dilC, 2) # n length
pad = tuple(map( lambda dil: (dil+1,dil+1), dilC )) # SAME padding with dilation # n length
stride = (1,)*2 # n length
return lax.conv_general_dilated( X.astype(jnp.float32) , W.astype(jnp.float32) , padding = pad, lhs_dilation = dilC, window_strides = stride)
# transposed kernel on non batch dims
WJ = jnp.flip( W_, axis = range(0+2,2+2) )# from 2 to n+2
out = conv_dilate( Y, WJ, dilC= dil)
out.shape # (1, 1, 28, 28)
import torch
import torch.nn.functional as F
Y = torch.tensor(Y)
W = torch.tensor( WJ.swapaxes(0,1)) #wherether I feed W_ or WJ, will change a lot the L2 norm of the output (the transpose of spatial dimensions should be done internally in conv_transpose2d)
K = F.conv_transpose2d( input=Y, weight=W ,stride = (2,2) ) # dilation parameter is for kernel_dilation here
K_ = torch.tensor(out )
print(torch.allclose(K, K_))I checked and also confirm the result does not correspond. |
Beta Was this translation helpful? Give feedback.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
-
Description
I am new to Jax... I am facing this issue, the Jax based conv_transpose result does not match with pytorch and I am certain that the pytorch result is correct.
When I call this function with 👇
input.shape = (1, 15, 13, 13) description (batch, channels, H, W)
W.shape = (15, 1, 4, 4) description (out_channel, in_channel, Kh, Kw)
stride = (2,2)
padding='VALID'
I should get an output with shape (1, 1, 28, 28), instead I get the following error
If I do not do the transpose, I am not getting any errors but the result is wrong.
My understanding is the rhs dimension number 'IOHW' (in this case) must match as following
I -> input_channel
O -> out_channel
H -> height
W -> width
this is the reason I did the transpose, so after the transpose the weight matrix W is of shape (1, 15, 4, 4) instead of (15, 1, 4, 4)
I also tried the following.
This also threw the same error.
what am I missing?
System info (python version, jaxlib version, accelerator, etc.)
Beta Was this translation helpful? Give feedback.
All reactions