perf: speed up mu.py

Apply flatten strategy. Reduce all zeros and compute on dense 1-D array instead of sparse 4-D tensor.