def data_to_stats(ob, z):
    """
    pointwise sufficient statstics
    stat1 : sum of I[z_n=k], S * B * K * 1
    stat2 : sum of I[z_n=k]*x_n, S * B * K * D
    stat3 : sum of I[z_n=k]*x_n^2, S * B * K * D
    """
    stat1 = z.sum(2).unsqueeze(-1)
    z_expand = z.unsqueeze(-1).repeat(1, 1, 1, 1, ob.shape[-1])
    ob_expand = ob.unsqueeze(-1).repeat(1, 1, 1, 1, z.shape[-1]).transpose(-1, -2)
    stat2 = (z_expand * ob_expand).sum(2)
    stat3 = (z_expand * (ob_expand**2)).sum(2)
    return stat1, stat2, stat3