[bug] when dim=1

[Gabibing]

issue

When using ResidualVQ with feature dimension dim=1 and passing a 2D mask of shape [B, T], masking is not handled correctly.
When mask=None is used as a workaround, the output’s batch dimension collapse to 1.

vq = ResidualVQ(dim=1, num_quantizers=1, codebook_size=4)

error log

vector_quantize_pytorch.py", line 1349, in forward
[rank0]:     quantize = einx.where(
                        ^^^^^^^^^^^
einx.expr.stage3.SolveValueException: Failed to solve values of expressions.
Found contradictory values {5936, 1} for equivalent expressions {'d', '5936', '1'}
Input:
    'b n = 8 742'
    'b n d = 8 742 5936'
    'b n d = 8 742 1'
    'b n d = None'

[lucidrains]

@Gabibing hey, can you show me a short script that reproduces the error?

this runs fine for me

import torch
from vector_quantize_pytorch import ResidualVQ

residual_vq = ResidualVQ(
    dim = 1,
    num_quantizers = 2,
    codebook_size = 512
)

x = torch.randn(2, 16, 1)
mask = torch.ones((2, 16)).bool()

quantized, indices, commit_loss = residual_vq(x, mask = mask)

assert quantized.shape == x.shape

[Gabibing]

@lucidrains here is my code. my batch size was 8.

class ContentQuantizer(nn.Module):
    def __init__(self, in_dim=1024, hidden_dim=256, dir_codebook_size=2048, mag_codebook_size=4):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(in_dim),          # feature-dim LN
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.LayerNorm(hidden_dim),
        )

        self.dir_vq = ResidualVQ(
            dim=hidden_dim, num_quantizers=1, codebook_size=dir_codebook_size,
            stochastic_sample_codes=True, sample_codebook_temp=0.
        )
        self.mag_vq = ResidualVQ(
            dim=1, num_quantizers=1, codebook_size=mag_codebook_size,
        )

        self.K_dir = dir_codebook_size
        self.K_mag = mag_codebook_size

        self.register_buffer('eps', torch.tensor(1e-7, dtype=torch.float32))

    def forward(self, x, mask, sample_codebook_temp=0.):
        #x: [B, T, 1024]
        #mask: [B, T]
        z = self.net(x)    # [B,T,hidden_dim]

        gain = z.norm(dim=2, keepdim=True).clamp_min(self.eps)        # [B,T,1]
        z_dir = z / gain                                          # [B,T,H]
        log_gain = torch.log(gain)                                # [B,T,1]

        q_dir, idx_dir, commit_dir = self.dir_vq(
            z_dir, mask=mask, sample_codebook_temp=sample_codebook_temp
        )                                                         
        q_lg, idx_mag, commit_mag = self.mag_vq(
            log_gain, mask=mask 
        )

        q_gain = torch.exp(q_lg)                                       # [B,T,1]
        quantized = q_dir * q_gain                                     # [B,T,H]

        indices = idx_mag * self.K_dir + idx_dir                  # [B,T]

        commit_loss = commit_dir + commit_mag

        return quantized, indices, commit_loss

[lucidrains]

@Gabibing can you make sure it is a runnable script and produces the error?

[Gabibing]

@lucidrains

code

import torch
import torch.nn as nn
import torch.nn.functional as F


from vector_quantize_pytorch import ResidualVQ

class ContentQuantizer(nn.Module):
    def __init__(self, in_dim=1024, hidden_dim=256, codebook_size=2048):
        super().__init__()
        self.net = nn.Linear(in_dim, hidden_dim)
        self.quantizer = ResidualVQ(
            dim=hidden_dim,
            num_quantizers=1,
            codebook_size=codebook_size,
            stochastic_sample_codes=True,
            sample_codebook_temp=1.0,
        )

    def forward(self, x, mask, sample_codebook_temp=0.):  # x: [B,T,1024]
        x = self.net(x)    # [B,T,hidden_dim]
        print(x.shape)
        quantized, indices, commit_loss = self.quantizer(x, mask=mask, sample_codebook_temp=sample_codebook_temp)
        return quantized, indices, commit_loss

if __name__ == '__main__':
   vq = ContentQuantizer(1024, 1, 2048)

   x = torch.randn(8, 700, 1024)
   mask = torch.ones(8, 700, dtype=torch.bool)
   q, i, c = vq(x, mask)
   print(q.shape)

error

einx.expr.stage3.SolveValueException: Failed to solve values of expressions. Found contradictory values {5600, 1} for equivalent expressions {'d', '5600', '1'}
Input:
    'b n = 8 700'
    'b n d = 8 700 5600'
    'b n d = 8 700 1'
    'b n d = None'

[lucidrains]

@Gabibing thanks for reporting this, turns out it was an issue with a squeeze within the rotation trick functions

can you try it again?

[Gabibing]

@lucidrains It works perfectly. Thank you!

[lucidrains]

@lucidrains It works perfectly. Thank you!

music to my ears

best with the neural network you are training