def __init__(self, N, a, b):
        """Implements update for matrix A = -(L+aI) for forward, backward, bilinear, zoh discretizations.
        a: scalar, the element on the diagonal
        b: scalar, so that B = b * ones vector
        """

        # Can't wrap scalars with torch.Tensor(), while torch.tensor(a) gives double instead of float or something
        # super().__init__(N, {'a': [a], 'b': [b]}, **kwargs) # TODO this should register b and then construct self.B using a @property, like in Toeplitz (but is slightly slower in the non-learnable case)
        params = {
            'a': torch.tensor(a, dtype=torch.float),
            'b': torch.tensor(b, dtype=torch.float),
        }
        super().__init__(N, params)

        # self.N = N
        # self.a = a
        # self.b = b

        # self.register_buffer('A', self.construct_A())
        # self.register_buffer('B', b * torch.ones(N))

        # self.register_buffer('I', torch.eye(N))
        self.register_buffer('arange', torch.arange(N-1))
Implements update for matrix A = -(L+aI) for forward, backward, bilinear, zoh discretizations.
        a: scalar, the element on the diagonal
        b: scalar, so that B = b * ones vector
        
analyze/convnexts4nd/src/models/hippo/transition.py