Z_inferred = hidden['Z'].squeeze()
B_inferred = hidden['B'].squeeze()
import scipy
scipy.stats.norm(loc=np.matmul(Z_inferred, B_inferred.T),
                 scale=np.ones(X_inferred.shape)).logpdf(X)