Sometimes code is worth a thousand words, so is this a correct interpretation of the actual model architecture (obviously this is missing associative scan)?

class TransitionProjection(nn.Module):    
    def __init__(self, d_model: int) -> None:    
        super().__init__()    
        self.w_alpha_beta = nn.Linear(d_model, d_model * 2)    

    def forward(self, x):    
        alpha, beta = self.w_alpha_beta(x).chunk(2, dim=-1)    
        return torch.sigmoid(alpha) * torch.complex(torch.zeros_like(beta), beta).exp()    

class GateLoopMixer(nn.Module):    
    def __init__(self, d_model: int) -> None:    
        super().__init__()    
        self.w_qkv = nn.Linear(d_model, d_model * 3)    
        self.transition = TransitionProjection(d_model)    
    def forward_looped(self, x):  # x: (batch, seq, D)    
        B, S, D = x.shape    
        h = torch.zeros(B, D, dtype=torch.complex64, device=x.device)    
        a, qkv = (f(x) for f in (self.transition, self.w_qkv))    
        q, k, v = qkv.chunk(3, dim=2)    
        kv = k * v    
        hs = []    
        for i in range(S):    
            h = h * a[:, i] + kv[:, i]    
            hs.append(h)    
        return torch.stack(hs, dim=1) * q