def precompute_forward(self, delta):
        """Store elements along the diagonals of (I + d A)."""

        if isinstance(delta, float):
            delta = torch.tensor(delta).to(self.I)
        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1)
        a_ = 1. - delta * self.a # (..., 1)
        if self.N == 1:
            return a_
        return torch.cat((a_, -delta*delta.new_ones(self.N-1)), -1) # (..., N)
Store elements along the diagonals of (I + d A).
analyze/convnexts4nd/src/models/hippo/transition.py