import torch


def sdd_bmm_diag_torch(t1, t2):
    """
    Perform bmm and diagonal for sparse x dense -> dense. The diagonalized result is returned in vector tensor.
    With s_t1.shape = (b, x, s), d_t2.shape = (b, s, x), the output shape is (b, x).
    This method avoids a temporal (b, x, x) for memory efficiency.
    :param t1: tensor 1
    :param t2: tensor 2
    :return: bmm_diag result in dense
    """
    assert t1.shape[0] == t2.shape[0], 'Batch size mismatch.'
    assert t1.shape[2] == t2.shape[1] and t1.shape[1] == t2.shape[2], 'Matrix shape mismatch.'
    if t1.is_sparse:
        d_t1 = t1.transpose(1, 2).to_dense()
        outp = torch.sum(d_t1.mul_(t2), dim=1)
    else:
        d_t2 = t2.transpose(1, 2).to_dense()
        outp = torch.sum(d_t2.mul_(t1), dim=2)
    return outp