`get_weight` and `recompute_head` refactor

[rodrigo-o]

This is a known issue as stated here: https://www.youtube.com/live/8_wCUL9T-xQ?si=OIVdEIGV8xTLstyl&t=1867

It doesn't appear to affect performance after #1380, but it may be something to take into consideration in the future