https://github.com/cyh24/Joint-Bayesian/blob/fa38b7a5371652318e58577717e8fb45bdca88ad/src/joint_bayesian.py#L81 The final `reshape` operation seems to be unnecessary.