    def forward(self, x, mask=None):
        """
        Forward function of window attention 3d.
        Args:
            x (`tensor`):
                Input features with shape of (num_windows*b, n, c). See ref2, n = (t * h * w) / (p * m * m).
                num_windows is Wd * Wh * Ww.
                b is batch size. c is dimension or channels size, same as dim param in init function.
            mask (`tensor`):
                Mask with shape of (num_windows, n, n) or None
        """
        b, n, c = x.shape
        qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # B_, nH, N, C

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index[:n, :n].reshape(-1)]\
            .reshape(
            n, n, -1)  # Wd*Wh*Ww,Wd*Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0) # B_, nH, N, N

        if mask is not None:
            nw = mask.shape[0]
            attn = attn.view(b // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, n, n)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(b, n, c)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
Forward function of window attention 3d.
Args:
x (`tensor`):
Input features with shape of (num_windows*b, n, c). See ref2, n = (t * h * w) / (p * m * m).
num_windows is Wd * Wh * Ww.
b is batch size. c is dimension or channels size, same as dim param in init function.
mask (`tensor`):
Mask with shape of (num_windows, n, n) or None
towhee/models/layers/window_attention3d.py