The Jax primitive jax.lax.conv_transpose result does not match the result of torch.nn.functional.conv_transpose2d

[SanjithKumar2]

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.

def _conv_transpose(input, W, stride, padding):
    return jax.lax.conv_transpose(
        lhs = input,
        rhs = jnp.transpose(W,(1,0,2,3)),
        strides = stride,
        padding=padding,
        dimension_numbers=('NCHW','IOHW','NCHW')
    )

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

ValueError                                Traceback (most recent call last)
Cell In[56], [line 8](vscode-notebook-cell:?execution_count=56&line=8)
      [6](vscode-notebook-cell:?execution_count=56&line=6) z_ = _conv_forward(X,W,(2,2),'VALID') + 1e-12
      [7](vscode-notebook-cell:?execution_count=56&line=7) S = R_/z_
----> [8](vscode-notebook-cell:?execution_count=56&line=8) S_w = _conv_transpose(S,W,(2,2),'VALID')
      [9](vscode-notebook-cell:?execution_count=56&line=9) R_in = X * S_w
.....
File jax\_src\lax\convolution.py:374, in _conv_general_dilated_shape_rule(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, **unused_kwargs)
    [370](/jax/_src/lax/convolution.py:370) if not core.definitely_equal(quot, rhs.shape[dimension_numbers.rhs_spec[1]]):
    [371](/jax/_src/lax/convolution.py:371)   msg = ("conv_general_dilated lhs feature dimension size divided by "
    [372](/jax/_src/lax/convolution.py:372)          "feature_group_count must equal the rhs input feature dimension "
    [373](/jax/_src/lax/convolution.py:373)          "size, but {} // {} != {}.")
--> [374](/jax/_src/lax/convolution.py:374)   raise ValueError(msg.format(lhs_feature_count, feature_group_count,
    [375](/jax/_src/lax/convolution.py:375)                               rhs.shape[dimension_numbers.rhs_spec[1]]))
    [376](/jax/_src/lax/convolution.py:376) if rhs.shape[dimension_numbers.rhs_spec[0]] % feature_group_count:
    [377](/jax/_src/lax/convolution.py:377)   msg = ("conv_general_dilated rhs output feature dimension size must be a "
    [378](/jax/_src/lax/convolution.py:378)          "multiple of feature_group_count, but {} is not a multiple of {}.")

ValueError: conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but 15 // 1 != 1.

If I do not do the transpose, I am not getting any errors but the result is wrong.

def _conv_transpose(input, W, stride, padding):
     return jax.lax.conv_transpose(
         lhs = input,
         rhs = W,
         strides = stride,
         padding=padding,
         dimension_numbers=('NCHW','IOHW','NCHW')
    )

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.

def _conv_transpose(input, W, stride, padding):
     return jax.lax.conv_transpose(
         lhs = input,
         rhs = W,
         strides = stride,
         padding=padding,
         dimension_numbers=('NCHW','OIHW','NCHW')
    )

This also threw the same error.

what am I missing?

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.7.2
jaxlib: 0.7.2
numpy:  2.3.3
python: 3.12.9 | packaged by Anaconda, Inc. | (main, Feb  6 2025, 18:49:16) [MSC v.1929 64 bit (AMD64)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Windows', node='LAPTOP-1K8EP6A1', release='11', version='10.0.26100', machine='AMD64')

$ nvidia-smi
Mon Oct 13 19:07:41 2025       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 581.42                 Driver Version: 581.42         CUDA Version: 13.0     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4060 ...  WDDM  |   00000000:01:00.0 Off |                  N/A |
| N/A   41C    P8              1W /  140W |       0MiB /   8188MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
...
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

[hawkinsp]

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 Solution

You 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:

import jax
import jax.numpy as jnp

def _conv_transpose(input_tensor, W, stride, padding):
    # Do NOT transpose W.
    # The 'IOHW' spec correctly interprets the shape (15, 1, 4, 4)
    # as Input=15 channels and Output=1 channel for the kernel.
    return jax.lax.conv_transpose(
        lhs=input_tensor,
        rhs=W,
        strides=stride,
        padding=padding,
        dimension_numbers=('NCHW', 'IOHW', 'NCHW')
    )

# --- Example usage with your shapes ---
# Input tensor with 15 channels
S = jnp.ones((1, 15, 13, 13))

# Kernel for a forward convolution that takes 1 channel and produces 15.
# For the transpose op, this means it takes 15 channels and produces 1.
W = jnp.ones((15, 1, 4, 4))

# This will now work without errors
output = _conv_transpose(S, W, (2, 2), 'VALID')

print(f"Output shape: {output.shape}")
# Expected output: Output shape: (1, 1, 28, 28)

---

Did that answer your question?

[DiagRisker]

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

[SanjithKumar2]

Thank you @DiagRisker next time will make sure that!

[SanjithKumar2]

@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)
W -> (15, 1, 4, 4) (the weight matrix) (no of filters(out_channel), in_channel, kH, kW)
stride -> (2, 2)
Z, Z_-> (1, 15, 13, 13) (the output after performing convolution operation) (batch_size, out_channel, oH, oK)

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 True

also 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)
W -> (15, 1, 4, 4) (the weight matrix from the forward pass) (no of filters(out_channel (according to the forward pass)), in_channel, kH, kW)

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,
N, C, H, W -> batch size, lhs channels, height, width (1, 15, 13, 13)
O, I, H, W -> no_of_filters i.e. out_channel, in_channel, kernel height, kernel width (15, 1, 4, 4)
N, C, H, W -> batch size, output channels, height, width (1, 1, 28, 28)

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 False

When 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...

[hawkinsp]

Please provide a minimal, self-contained, runnable example of the problem. i.e, something I can run.

[SanjithKumar2]

@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_))

[DiagRisker]

@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)
jax.lax.conv_transpose has also rhs_dilation parameter, but it calls also conv_general_dilated, with a slight change in how W and padding are picked. cf : https://github.com/jax-ml/jax/blob/main/jax/_src/lax/convolution.py#L285-L361

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_b

It 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)

[SanjithKumar2]

@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

[DiagRisker]

@SanjithKumar2
Clarification: the input dilation, kernel_dilation in lax.conv_general_dilated do not have the same definition as ConvTranspose from Pytorch, you have to swap stride and dilation parameter (cf: https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md)

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.
The padding should be good, however reading the torch git does not make it easy to see exactly what prepocessing is done internally on torch.nn.functional.conv_transpose2d torch side (couldn't find the source code, could be a direct bind in C++, but not found at 3rd glance)

[SanjithKumar2]

@DiagRisker I have found the answer... it took time I went through the Jax implementation of conv_transpose and I found out that Jax doesn't automatically flip the kernel by 90 degrees i.e. swap axes which is what conv_transpose2d does in pytorch (like in gradient calculation), just by setting transpose_kernel=True in the jax.lax.conv_transpose call the kernel is flipped and the result matches the pytorch implementation.

# JAX implementation
# L.conv_transpose(lhs, rhs, stride, padding, dimension_numbers=dimension_numbers, transpose_kernel=True)
    
 Array([[[[-1.2976712 ,  7.6736045 ,  3.9103894 , ...,  4.156144  ,
          -3.4796333 , -0.91263664],
         [ 4.741432  ,  3.938364  ,  7.454193  , ..., -1.333845  ,
           2.9775846 , -2.5136209 ],
         [ 1.3718653 , -1.9034076 ,  5.061095  , ..., 10.3180685 ,
          -1.7714388 ,  0.5106771 ],
         ...,
         ...,
         [ 2.7572577 ,  6.9018383 , -9.385033  , ..., -2.942164  ,
          12.533854  , -3.793242  ],
         [ 4.2677245 ,  3.0173094 ,  2.7786632 , ..., -3.1466503 ,
           5.4446726 ,  2.2346032 ],
         [-3.792156  , -4.6612625 , -1.7187886 , ...,  1.1125679 ,
           6.20341   , -2.4789004 ]]]], dtype=float32)
           
# Torch implementation
# F.conv_transpose2d(lhs_t, rhs_t,stride=(2,2))

tensor([[[[-1.2977,  7.6736,  3.9104,  ...,  4.1561, -3.4796, -0.9126],
          [ 4.7414,  3.9384,  7.4542,  ..., -1.3338,  2.9776, -2.5136],
          [ 1.3719, -1.9034,  5.0611,  ..., 10.3181, -1.7714,  0.5107],
          ...,
          ...,
          [ 2.7573,  6.9018, -9.3850,  ..., -2.9422, 12.5339, -3.7932],
          [ 4.2677,  3.0173,  2.7787,  ..., -3.1466,  5.4447,  2.2346],
          [-3.7922, -4.6613, -1.7188,  ...,  1.1126,  6.2034, -2.4789]]]])

my mistake not looking at the parameter list properly, but ended up learning a lot about how Jax works... Thank you @hawkinsp @DiagRisker