import torch


def bdiag_prod(M1, M2):
    """Computes matrix-matrix product of
    two block matrices containing only 
    diagonal submatrices.
    Input shapes: (N2, P, N1), (N1, P, N3)
    Output shape: (N2, P, N3)"""

    assert M1.size(2)==M2.size(0), "Dimension 2 of M1 does not match dimension 0 of M2"
    assert len(M1.size())==3 and len(M2.size())==3, "Expected input tensors to be of shape (N, P, N)"

    M = torch.einsum('ijk,kjl->ijl', M1, M2)

    return M