Update entropy_pmf to sum over last axis only.

This let's you calculate multiple entropies at once.