Add tests for MPS accelerator mixed precision device selection

Author: bhimrazy

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -1084,3 +1084,25 @@ def test_precision_selection_model_parallel(precision, raises, mps_count_0):
     error_context = pytest.raises(ValueError, match=f"does not support .*{precision}") if raises else nullcontext()
     with error_context:
         _AcceleratorConnector(precision=precision, strategy=ModelParallelStrategy())
+
+
+@RunIf(mps=True)
+@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
+def test_mps_amp_device_selection(precision):
+    """Test that MPS accelerator with mixed precision correctly sets device to 'mps' instead of 'cuda'."""
+    connector = _AcceleratorConnector(accelerator="mps", precision=precision)
+    assert isinstance(connector.precision_plugin, MixedPrecision)
+    # Verify the device parameter is set to "mps" for MPS accelerator
+    assert connector.precision_plugin.device == "mps"
+
+
+@RunIf(mps=True)
+def test_mps_amp_device_selection_vs_cpu():
+    """Test that MPS AMP device selection differs from CPU AMP device selection."""
+    # MPS with mixed precision should use "mps" device
+    mps_connector = _AcceleratorConnector(accelerator="mps", precision="bf16-mixed")
+    assert mps_connector.precision_plugin.device == "mps"
+
+    # CPU with mixed precision should use "cpu" device
+    cpu_connector = _AcceleratorConnector(accelerator="cpu", precision="bf16-mixed")
+    assert cpu_connector.precision_plugin.device == "cpu"