def __init__(
            self,
            d,
            d_model=-1, # overloading this term, same as memory_order or N
            measure='legs', # 'legs', 'legt' main ones; can also try 'lagt'
            measure_args={},
            learn=0, # 0 means no learn, 1 means same A matrix for each hidden feature H, 2 means different A matrix per feature. 1 does not change parameter count. 2 adds parameters but does not slow down
            lr=0.0001, # controls learning rate of transition parameters
            noise=0.0, # injects input noise to the state space system
            init='normal', # for debugging, but might be useful?
            dt=None,
            channels=1, # denoted by M below
            bias=False,
            activation='gelu',
            ff=True,
            weight_norm=False,
            dropout=0.0,
            l_max=-1,
        ):
        """
        N: the order of the HiPPO projection
        dt: discretization step size - should be roughly inverse to the length of the sequence
        """
        super().__init__()
        self.d = d
        self.N = d_model if d_model > 0 else d
        self.dt = DictConfig({
            'min' : 0.001,
            'max' : 0.1,
            'learn' : False,
            'lr': 0.001,
            'init' : 'random',
        })
        if dt is not None: self.dt.update(dt)
        self.ff = ff
        self.bias = bias


        # Construct transition
        self.learn = learn
        if self.learn == 0:
            if measure == 'identity': # for testing
                A, B = torch.eye(self.N), torch.ones(self.N)
                self.transition = transition.ManualAdaptiveTransition(self.N, A, B)
            elif measure == 'random':
                A = torch.randn(self.N, self.N) / self.N # E[AA^T] = (1/N)I -- empirically I nans out
                B = torch.ones(self.N) # based on HiPPO matrices; worth trying random, haven't tried
                self.transition = transition.ManualAdaptiveTransition(self.N, A, B)
            elif measure == 'legt':
                # self.transition = transition.LegTAdaptiveTransition(self.N)
                self.transition = transition.LegTTriDInverseAdaptiveTransition(self.N, **measure_args)
            elif measure == 'cheb':
                self.transition = transition.ChebITriDInverseAdaptiveTransition(self.N, **measure_args)
            elif measure == 'chebii':
                self.transition = transition.ChebIITriDInverseAdaptiveTransition(self.N, **measure_args)
            elif measure == 'lagt':
                self.transition = transition.LagTCumsumAdaptiveTransition(self.N, **measure_args)
            elif measure == 'glagt':
                self.transition = transition.GLagTToeplitzAdaptiveTransition(self.N, **measure_args)
            elif measure == 'legs':
                self.transition = transition.LegSTriDInverseAdaptiveTransition(self.N, **measure_args)
            elif measure == 'jac':
                self.transition = transition.JacTriDInverseAdaptiveTransition(self.N, **measure_args)
            else:
                raise NotImplementedError
        elif self.learn == 1 or self.learn == 2:
            kwargs = {'trainable': True, 'lr': lr}
            kwargs.update(measure_args)
            if self.learn == 2:
                kwargs['batch'] = (self.d,)
            if measure == 'random':
                A = torch.randn(self.N, self.N) / self.N # E[AA^T] = (1/N)I . empirically I doesn't work, dunno why
                B = torch.ones(self.N) # based on HiPPO matrices; worth trying random, haven't tried
                self.transition = transition.ManualAdaptiveTransition(self.N, A, B, **kwargs)
            elif measure == 'legt':
                self.transition = transition.LegTTriDInverseAdaptiveTransition(self.N, **kwargs)
            elif measure == 'lagt':
                self.transition = transition.LagTTriDInverseAdaptiveTransition(self.N, **kwargs)
            elif measure == 'legs':
                self.transition = transition.LegSTriDInverseAdaptiveTransition(self.N, **kwargs)
            elif measure == 'cheb':
                self.transition = transition.ChebITriDInverseAdaptiveTransition(self.N, **kwargs)
            elif measure == 'chebii':
                self.transition = transition.ChebIITriDInverseAdaptiveTransition(self.N, **kwargs)
            elif measure == 'toep':
                self.transition = transition.LagTToeplitzAdaptiveTransition(self.N, **kwargs)
            else: raise NotImplementedError
        elif self.learn == 3: # for debugging
            A, B = hippo.transition(measure, self.N)
            B = B[:, 0]
            self.transition = transition.ManualAdaptiveTransition(self.N, A, B, trainable=True, lr=lr)
        else:
            raise NotImplementedError


        self.m = channels

        if init == 'normal':
            self.C = nn.Parameter(torch.randn(self.d, self.m, self.N))
            self.D = nn.Parameter(torch.randn(self.d, self.m))
        elif init == 'constant':
            self.C = nn.Parameter(torch.ones(self.d, self.m, self.N))
            self.D = nn.Parameter(torch.ones(self.d, self.m))
        elif init == 'uniform':
            self.C = nn.Parameter(1.732 * torch.rand(self.d, self.m, self.N))
            self.D = nn.Parameter(torch.randn(self.d, self.m))
        else: raise NotImplementedError
        if self.bias:
            self.E = nn.Parameter(torch.zeros(self.d, self.m))

        if self.dt.init == 'uniform':
            log_dt = torch.linspace(math.log(self.dt.min), math.log(self.dt.max), self.d)
        elif self.dt.init == 'random':
            log_dt = torch.rand(self.d) * (math.log(self.dt.max)-math.log(self.dt.min)) + math.log(self.dt.min)
        else: raise NotImplementedError

        if self.dt.learn:
            self.log_dt = nn.Parameter(log_dt) # (H)
            self.log_dt._lr = self.dt.lr # register the parameter for the optimizer to reduce lr
        else:
            self.register_buffer('log_dt', log_dt)
        self.k = None
        self.noise = noise

        self.activate = Activation(activation)
        self.drop = nn.Dropout(dropout)

        if self.ff:
            self.output_linear = nn.Linear(self.m * self.d, self.d)

            if weight_norm:
                self.output_linear = nn.utils.weight_norm(self.output_linear)

        # For test time shift
        self.l_max = l_max
        self.last_len = -1

        N: the order of the HiPPO projection
        dt: discretization step size - should be roughly inverse to the length of the sequence
        
analyze/convnexts4nd/src/models/sequence/modules/lssl.py