def forward(self, u, *args, state=None, **kwargs):
        """
        u: (L, B, H) [21-09-29] Our backbone now passes inputs as (B, L, H). This calss originally expected (L, B, H) so we transpose accordingly
        state: (B, H, N) previous hidden state of the recurrence
        """
        next_state = None

        u = u.transpose(0, 1)

        # Construct dt (H)
        dt = torch.exp(self.log_dt) # Note: if dt is not learnable this slightly wastes computation, but it isn't a bottleneck

        ## # Calculate test-time shift
        # changed sampling rate; uncache Krylov
        if self.last_len != u.shape[0]:
            self.k = None
            self.last_len = u.shape[0]
        # Calculate change from train sampling rate
        if self.l_max > 0:
            rate = self.l_max / u.shape[0]
            # if rate != 1.0: dt = dt * rate
            if rate != 1.0: rate = round(rate)
            else: rate = None
        else:
            rate = None


        # We need to compute the "recurrence" if
        # (*) there is noise or an initial state
        # (*) we're learning the system A, B
        # (*) first pass
        kb = [] # will store the B vectors for Krylov computation
        _learn = (self.dt.learn or self.learn) and self.training # need to learn and it's training time # TODO this ignores the last training minibatch if no test time shift (prev batch's K gets cached)... should recalculate A in the last_len check ideally
        _conv = _learn or self.k is None or u.shape[0] > self.k.shape[-1] # or rate
        _noise = self.noise > 0.0 and self.training
        if _conv:
            B = self.transition.gbt_B(dt) # (..., N) depending if learn=2
            kb.append(B)
        if _noise:
            noise = self.noise * torch.randn(self.d, self.N, dtype=u.dtype, device=u.device) # (H, N)
            kb.append(noise)

        A = None
        if len(kb) > 0:
            if rate is not None:
                dt = dt * rate

            A = self.transition.gbt_A(dt) # (..., N, N) (..., N)

            # Adjust by rate
            # if _conv and rate is not None:
            #     while rate > 1:
            #         B = B + torch.sum(A * B.unsqueeze(-2), dim=-1) # (I + A) @ B
            #         A = A @ A
            #         rate //= 2

            kb = [b.broadcast_to(dt.shape+(self.N,)) for b in kb]
            kb = torch.stack(torch.broadcast_tensors(*kb), dim=0) # each (..., N)
            krylovs = krylov(u.shape[0], A, kb) # (H, N, L) each
            k_noise, k_conv = torch.split(
                krylovs,
                split_size_or_sections=[int(_noise), int(_conv)],
                dim=0
            )
            if _conv: # Cache the Krylov matrix K(A, B)
                self.k = k_conv.squeeze(0) # (H, N, L)
            if _noise:
                k_noise = k_noise.squeeze(0) # (H, N, L)

        # Convolution
        y = linear_system_from_krylov(u, self.C, self.D, self.k[..., :u.shape[0]]) # (L, B, H, M)
        if _noise:
            k_noise = torch.cumsum(k_noise, dim=-1) # (H, N, L) w + Aw + A^2w + ...
            k_noise = contract('h m n, h n l -> l h m', self.C, k_noise) # C @ k
            y = y + k_noise.unsqueeze(1) # (L, B, H, M)
            y = y + self.noise * torch.randn(y.shape, dtype=u.dtype, device=u.device)

        # State needs a special case because it has a batch dimension
        if state is not None: # (B, H, N)
            if A is None: A = self.transition.gbt_A(dt) # (..., N, N) (..., N)

            ATC, ATL = krylov(u.shape[0], A.transpose(-1,-2), self.C.transpose(0, 1), return_power=True) # (M, H, N, L), (H, N, N) represents A^T C and (A^T)^L
            y = y + contract('mhnl, bhn -> lbhm', ATC, state)

            # Compute next state
            with torch.no_grad():
                next_state = contract('hnp, bhn -> bhp', ATL, state)
                if _noise:
                    next_state = next_state + k_noise[..., -1]
                next_state = next_state + contract('lbh, hnl -> bhn', u.flip(0), self.k[:..., u.shape[0]]) # (B, H, N)
                next_state = contract('hnp, bhp -> bhn', A, next_state)
                next_state = next_state.detach() # TODO necessary?

            # Debugging code useful for checking if state computation is correct
            # from models.functional.unroll import variable_unroll_sequential, variable_unroll
            # B = self.transition.gbt_B(dt)
            # inps = B*u.unsqueeze(-1) # (L, B, H, N)
            # inps[0] = inps[0] + state
            # xx = variable_unroll(A, inps, variable=False)
            # yy = torch.sum(self.C * xx.unsqueeze(-2), dim=-1)
            # yy = yy + u.unsqueeze(-1) * self.D # true output y; should equal y
            # xx_ = variable_unroll(A, B*u.unsqueeze(-1), variable=False)
            # yy_ = torch.sum(self.C * xx_.unsqueeze(-2), dim=-1)
            # yy_ = yy_ + u.unsqueeze(-1) * self.D # output without state; should equal y before the C A^T S term was added
            # ss = (A @ xx[-1].unsqueeze(-1)).squeeze(-1) # should equal next_state
            # breakpoint()
            # y = z

        # bias term
        if self.bias:
            y = y + self.E

        y = self.drop(self.activate(y))

        y = rearrange(y, 'l b h m -> l b (h m)') # (L, B, H*M)

        if self.ff:
            y = self.output_linear(y) # (L, B, H)
        y = y.transpose(0, 1) # Back to (B, L, H) as expected
        return y, next_state

        u: (L, B, H) [21-09-29] Our backbone now passes inputs as (B, L, H). This calss originally expected (L, B, H) so we transpose accordingly
        state: (B, H, N) previous hidden state of the recurrence
        
analyze/convnexts4nd/src/models/sequence/modules/lssl.py