Forward function. Args: input (torch.Tensor): Input tensor. First dimension should be the batch size Returns:

state: (..., s, N) extra tensor that augments B rate: sampling rate factor returns: (..., c+s, L)

u: (B H L) if self.transposed else (B L H) state: (H N) never needed unless you know what you're doing Returns: same shape as

inputs : (length, ...) output : (length, ..., N) where N is the order of the HiPPO projection

Compuates diagonal_mm of t1 and t2. args: t1: torch.Tensor = (batch_size, seq_len, num_attention_heads, hidden_size|number_of_diagona