Note: index_select(out=) optimization removed as it caused color polarization on MPS; using simple indexing instead