def gbt_A(self, dt, alpha=.5):
        """Compute the transition matrices associated with bilinear transform.

        dt: (...) broadcastable with self.batch_shape
        returns: (..., N, N)
        """

        # Solve (N, ...) parallel problems of size N
        dims = max(len(dt.shape), len(self.batch))
        I = self.I.view([self.N] + [1]*dims + [self.N])
        A = self.bilinear(dt, I, dt.new_zeros(*dt.shape), alpha=alpha) # (N, ..., N)
        A = rearrange(A, 'n ... m -> ... m n', n=self.N, m=self.N)
        return A
Compute the transition matrices associated with bilinear transform.

        dt: (...) broadcastable with self.batch_shape
        returns: (..., N, N)
        
analyze/convnexts4nd/src/models/hippo/transition.py