def __init__(self, N, params, trainable=False, lr=1.0, batch=()):
        """
        params: dict of Tensors that encode the parameters of the state system A, B.
        """

        super().__init__()
        self.N = N
        self.trainable = trainable
        self.batch = batch

        if self.trainable:
            for name, p in params.items():
                p = p.repeat(*batch, *[1]*len(p.shape))
                self.register_parameter(name, nn.Parameter(p))
                getattr(self, name)._lr = lr
        else:
            assert batch == (), "If not learnable, Transition should not have a batch dimension"
            for name, p in params.items():
                self.register_buffer(name, p)

        # Register some common buffers
        # (helps make sure every subclass has access to them on the right device)
        I = torch.eye(N)
        self.register_buffer('I', I)
        self.register_buffer('ones', torch.ones(N))
        self.register_buffer('arange', torch.arange(N))

        params: dict of Tensors that encode the parameters of the state system A, B.
        
analyze/convnexts4nd/src/models/hippo/transition.py