def precompute_backward(self, delta): # TODO should be called inverse?
        """Store elements along the diagonals of (I - d A)^{-1}.

        # a' = a + 1/dt
        delta: (...)
        output: (..., N)
        """

        if isinstance(delta, float):
            delta = torch.tensor(delta).to(self.I)
        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1)

        if self.N == 1:
            return 1. / (1. + self.a*delta + delta)

        ad = self.a*delta # (..., 1)
        ad_p1 = 1 + ad
        denom = ad_p1 + delta # 1 + a'
        denom_inv = denom.reciprocal() # 1. / denom
        s = - delta * denom_inv * denom_inv # -1/(1+a')^2
        b = ad_p1 * denom_inv # a' / (1 + a')
        pows = b ** self.arange ## TODO benchmark against cumprod or cumsum in log space
        tail = s * pows
        ret = torch.cat((denom_inv, tail), -1)
        return ret
        # ad = self.a*delta # (..., 1)
        # denom = 1 + ad + delta
        # s = - delta / denom# -1/(1+a')
        # b = (1 + ad) / denom # a' / (1 + a')
        # # pows = b ** torch.arange(self.N-1).to(self.I) ## TODO benchmark against cumprod or cumsum in log space
        # pows = b ** self.arange ## TODO benchmark against cumprod or cumsum in log space
        # tail = s * pows
        # ret = torch.cat((tail.new_ones(tail.shape[:-1]+(1,)), tail), -1)
        # ret = ret / denom
        # return ret
Store elements along the diagonals of (I - d A)^{-1}.

        # a' = a + 1/dt
        delta: (...)
        output: (..., N)
        
analyze/convnexts4nd/src/models/hippo/transition.py