def bmul(vec, mat, axis=0):
    """Expand vector for batchwise matrix multiplication.
    Parameters
    ----------
    vec : 2dtensor
        vector for multiplication
    mat : 3dtensor
        matrix for multiplication
    axis : int, optional
        batch axis, by default 0
    Returns
    -------
    3dtensor
        Product of matrix multiplication. (bs, n, m)
    """
    mat = mat.transpose(axis, -1)
    return (mat * vec.expand_as(mat)).transpose(axis, -1)