Method forward

bitnet/bit_mamba.py:20–25

(self, x)

        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        output = (
            x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
        )

        return output


class PScan(torch.autograd.Function):