import torch


def qzCx_sample(params):
    """
    Sample from posterior q(z|x).
    Assume normal with diagonal covariance.  
    Params: [B, LD, 2]
    Expect mu = params[..., 0]
           logvar = params[..., 1]
    """

    mu = params.select(-1, 0)
    logvar = params.select(-1, 1)
    sigma_sqrt = torch.exp(0.5 * logvar)
    epsilon = torch.randn_like(sigma_sqrt)
    return mu + sigma_sqrt * epsilon