For numerical code I like einops.reduce more than numpy/pytorch sum reductions because you can reduce over named dimensions. It’s much more readable than having to reason through axis indexing again every time you come back to the code