I’m training a perceiver transformer network and I’m trying to replace the explicitly added positional encoding with a positional encoding which is only added to the query and key vectors in the attention mechanism.
When I try to use this special attention I get the error in the title.
I tried stepping through backwards to find out where the error occurs but to no success as the backwards pass is abstrackted away in the c engine.
The relevant code is:
class FEPA_Attention(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., max_freq=512, freq_base=2, num_freq_bands=4):
        super().__init__()
        inner_dim = dim_head * heads
        self.cross = context_dim is not None
        context_dim = default(context_dim, query_dim)
        self.encoding_dim = ((num_freq_bands * 2) + 1)
        self.max_freq = max_freq
        self.num_freq_bands = num_freq_bands
        self.freq_base = freq_base

        self.query_dim = query_dim
        self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    @staticmethod
    def get_enc_pos(data, max_freq, num_freq_bands, freq_base):
        b, *axis, _, device = *data.shape, data.device
        axis_pos = list(map(lambda size: torch.linspace(-1., 1., steps=size, device=device), axis))

        pos = torch.stack(torch.meshgrid(*axis_pos), dim=-1)
        enc_pos = fourier_encode(pos, max_freq, num_freq_bands, freq_base)
        enc_pos = rearrange(enc_pos, '... n d -> ... (n d)')
        enc_pos = repeat(enc_pos, '... -> b ...', b=b)
        return enc_pos

    @staticmethod
    def concat_position_encoding(data, max_freq, num_freq_bands, freq_base):
        enc_pos = FEPA_Attention.get_enc_pos(data, max_freq, num_freq_bands, freq_base).detach().requires_grad_(False)
        data = torch.cat((data, enc_pos), dim=-1)
        data = rearrange(data, 'b ... d -> b (...) d')
        return data

    def forward(self, x, context=None, mask=None, get_attention_map=False):
        h = self.heads
        q = FEPA_Attention.concat_position_encoding(x, self.max_freq, self.num_freq_bands, self.freq_base)
        q = self.to_q(q)

        context = default(context, x)
        k = FEPA_Attention.concat_position_encoding(context, self.max_freq, self.num_freq_bands, self.freq_base)
        k = self.to_k(k)

        v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale

        if exists(mask):
            mask = rearrange(mask, 'b ... -> b (...)')
            max_neg_value = -torch.finfo(sim.dtype).max
            mask = repeat(mask, 'b j -> (b h) () j', h=h)
            sim.masked_fill_(~mask, max_neg_value)

        # attention, what we cannot get enough of
        attn = sim.softmax(dim=-1)

        out = einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        if get_attention_map:
            return self.to_out(out), attn
        return self.to_out(out)


def fourier_encode(x, max_freq, num_bands=4, base=2):
    x = x.unsqueeze(-1)
    device, dtype, orig_x = x.device, x.dtype, x

    scales = torch.logspace(1., log(max_freq / 2) / log(base), num_bands, base=base, device=device, dtype=dtype)
    scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]

    x = x * scales * pi
    x = torch.cat([x.sin(), x.cos()], dim=-1)
    x = torch.cat((x, orig_x), dim=-1)
    return x

How can I tackle this problem and find a solution?