Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py (Python)

Code before:
def test_bitsandbytes_precision_cuda_required(monkeypatch):
    monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
    monkeypatch.setitem(sys.modules, "bitsandbytes", Mock())
    with pytest.raises(RuntimeError, match="Bitsandbytes is only supported on CUDA GPUs"):
        _AcceleratorConnector(accelerator="cpu", plugins=BitsandbytesPrecision(mode="int8"))


@RunIf(min_torch="2.4")
@pytest.mark.parametrize(
    ("precision", "raises"),
    [("32-true", False), ("16-true", False), ("bf16-true", False), ("16-mixed", True), ("bf16-mixed", False)],
)
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(
    ("accelerator", "expected_device"),
    [
        ("mps", "mps"),
        ("cpu", "cpu"),
    ],
)
@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
def test_mps_amp_device_selection(accelerator, expected_device, precision):
    """Test that MPS accelerator with mixed precision correctly sets device to 'mps' instead of 'cuda'."""
    connector = _AcceleratorConnector(accelerator=accelerator, precision=precision)
    assert isinstance(connector.precision_plugin, MixedPrecision)
    assert connector.precision_plugin.device == expected_device


Review comment (suggestion):
```suggestion
@pytest.mark.parametrize("accelerator", ["mps", "mps"])
```
this part seems very redundant

Code after:
def test_bitsandbytes_precision_cuda_required(monkeypatch):
    monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
    monkeypatch.setitem(sys.modules, "bitsandbytes", Mock())
    with pytest.raises(RuntimeError, match="Bitsandbytes is only supported on CUDA GPUs"):
        _AcceleratorConnector(accelerator="cpu", plugins=BitsandbytesPrecision(mode="int8"))


@RunIf(min_torch="2.4")
@pytest.mark.parametrize(
    ("precision", "raises"),
    [("32-true", False), ("16-true", False), ("bf16-true", False), ("16-mixed", True), ("bf16-mixed", False)],
)
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("accelerator", ["mps", "cpu"])
@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
def test_mps_amp_device_selection(accelerator, precision):
    """Test that MPS accelerator with mixed precision correctly sets device to 'mps' instead of 'cuda'."""
    connector = _AcceleratorConnector(accelerator=accelerator, precision=precision)
    assert isinstance(connector.precision_plugin, MixedPrecision)
    assert connector.precision_plugin.device == accelerator