-p_X = p.sum(1)
-p_Z = p.sum(0)
-p_X_given_Z = p / p.sum(0, keepdim = True)
-p_Z_given_X = p / p.sum(1, keepdim = True)
-q_X_given_Z = q / q.sum(0, keepdim = True)
-q_Z_given_X = q / q.sum(1, keepdim = True)
+p_X = p_XZ.sum(1)
+p_Z = p_XZ.sum(0)
+p_X_given_Z = p_XZ / p_XZ.sum(0, keepdim = True)
+p_Z_given_X = p_XZ / p_XZ.sum(1, keepdim = True)