TypeError: dot_general requires contracting dimensions to have the same shape

[conceptofmind]

Hi,

I have recently been adapting some code from PyTorch to Jax. The code runs without conflict and executes properly in PyTorch but I am having difficulty properly implementing it in Jax. Any input would be greatly appreciated.

PyTorch:

import torch
from torch import nn

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

img = torch.randn(1, 3, 256, 256)

channels = 3
patch_height = 32
patch_width = 32
dim = 1024
patch_dim = channels * patch_height * patch_width

linear = nn.Sequential(
        Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
        nn.Linear(patch_dim, dim)
)

print(linear(img).shape)

Jax:

import jax
from jax import random
import jax.numpy as jnp
import einops
from einops import rearrange, repeat
from jax import vmap
import equinox as eqx
from equinox import Module

key = random.PRNGKey(0)

channels = 3
patch_height = 32
patch_width = 32
dim = 1024
patch_dim = channels * patch_height * patch_width

img = random.normal(key, (1, 3, 256, 256))

new_img = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width)

class MyLinear(eqx.Module):
    weight: jax.numpy.ndarray
    bias: jax.numpy.ndarray

    def __init__(self, in_dim, out_dim, *, key):
        wkey, bkey = random.split(key, 2)
        self.weight = random.normal(wkey, shape = (out_dim, in_dim))
        self.bias = random.normal(bkey, shape = (out_dim,))
    
    def __call__(self, x):
        return self.weight @ x + self.bias

linear = MyLinear(patch_dim, dim, key = key)

output = vmap(linear)(new_img)

print(output.shape)

Thank you,

Enrico

[YouJiacheng]

Hi. Your matrix multiplication shape is: (dim, patch_dim) @ (patch_num, patch_dim).
use
new_img = rearrange(img, 'b c (h p1) (w p2) -> b (p1 p2 c) (h w)', p1 = patch_height, p2 = patch_width)

[conceptofmind]

Hi @YouJiacheng ,

Thank you for the response. Unfortunately, when I reshape the array following the convention you provided I still receive another broadcasting error: ValueError: Incompatible shapes for broadcasting: ((1024, 64), (1, 1024)). In PyTorch, I believe the kernel has shape [outC, inC] and NumPy will broadcast from right to left. I may have a fundamental misunderstanding of how this works in Jax when compared to PyTorch. Do I need to reshape or transpose the dimensions again?

Thank you,

Enrico

[YouJiacheng]

Well, you need to change bias shape from (dim,) to (dim, 1)

[YouJiacheng]

In pytorch, linear layer is x @ weight.T + bias, where x.shape is (..., in_dim), weight.shape is (out_dim, in_dim), bias.shape is (out_dim,). But you use weight @ x + bias.

[conceptofmind]

Hi @YouJiacheng ,

I greatly appreciate your help. I was able to refactor the code per your specifications and get it to match the output of the PyTorch segment above:

key = random.PRNGKey(0)

channels = 3
patch_height = 32
patch_width = 32
dim = 1024
patch_dim = channels * patch_height * patch_width

img = random.normal(key, (1, 3, 256, 256))

class PyTorch_Linear(eqx.Module):
    weight: Array
    bias: Optional[Array]
    out_dim: int = static_field()
    in_dim: int = static_field()
    use_bias: bool = static_field()

    def __init__(self, in_dim, out_dim, use_bias = True, *, key):
        super().__init__()

        wkey, bkey = random.split(key, 2)

        self.out_dim = out_dim
        self.in_dim = in_dim
        self.use_bias = use_bias

        self.weight = random.normal(wkey, shape = (out_dim, in_dim))
        
        if self.use_bias:
            self.bias = random.normal(bkey, shape = (out_dim,))
        else:
            self.bias = None
    
    def __call__(self, x, *, key = None):
        x = x @ self.weight.T
        if self.bias is not None:
            x = x + self.bias
        return x

patch_embedding = eqx.nn.Sequential([
    Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
    PyTorch_Linear(patch_dim, dim, key = key)
])

output = patch_embedding(img)
print(output.shape)

I definitely had some misunderstanding of how to properly handle the input/output dimensions and array shapes.

Thank you again,

Enrico