def forward_mult(self, u, delta, transpose=False):
        """Computes (I + delta A) u.

        A: (n, n)
        u: (..., n)
        delta: (...) or scalar

        output: (..., n)
        """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1)

        if transpose:
            x = torch.cumsum(u.flip(-1), -1).flip(-1)
        else:
            x = torch.cumsum(u, -1)
        x = x + u * self.a
        x = u - delta * x # Because A is negated in the representation
        return x
Computes (I + delta A) u.

        A: (n, n)
        u: (..., n)
        delta: (...) or scalar

        output: (..., n)
        
analyze/convnexts4nd/src/models/hippo/transition.py