Simplify MPS support

[rasbt]

Simplified the MPS support based on your suggestion @t-vi (CC @Andrei-Aksionov).

It's a bit trickier to have unit tests for this behavior now (because we can't force to run the MPS code on non-Mac machines now), but it's fine because we know the previous unit tests for equivalency passed. And sure, we can have macOS tests for this alternative path, but the problem is that MPS has subtle numerical differences, so I don't expect the results of MPS and CPU to be equivalent. But we can address this is in #1725

Note that I also changed the default type for MPS to bf16 because it seems supported now. If we use regular float 16, we can get weird inference results for many models, which is why I changed it.

[Andrei-Aksionov]

Perhaps it's possible to mock tensor's device output, so you can test both code execution paths.

[t-vi]

It's a bit trickier to have unit tests for this behavior now (because we can't force to run the MPS code on non-Mac machines now),

what you can do is

• make the mps variant a separate function with the same signature (called with if on the device type),
• run the model with / without mocking the batched_index_copy to the mps variant.