Revert "Revert MPS changes (#78335)"
This reverts commit ffb31014847faccb724ef70924ba936b75144c85.
Reverted https://github.com/pytorch/pytorch/pull/78335 on behalf of https://github.com/malfet due to as it messes with the change history and does not fix iOS build issues
