def dist_forward(distribution, x):
    """
    Forward pass with an arbitrary PyTorch distribution.

    Args:
        distribution: PyTorch base distribution which is used to compute the log probabilities of x.
        x: Input to compute the log probabilities of.
           Shape [n, d].

    Returns:
        torch.Tensor: Log probabilities for each feature.
    """
    # Make room for multiplicity of layer
    # Output shape: [n, d, 1]
    x = x.unsqueeze(2)

    # Compute gaussians
    # Output shape: [n, d, multiplicity]
    x = distribution.log_prob(x)

    return x