def test_xla_fsdp_setup_optimizer_validation():
    """Test that `setup_optimizer()` validates the param groups and reference to FSDP parameters."""
    module = nn.Linear(2, 2)
    strategy = XLAFSDPStrategy(
        parallel_devices=XLAAccelerator.get_parallel_devices(XLAAccelerator.auto_device_count()),
    )
    bad_optimizer = Adam(module.parameters())
    with pytest.raises(ValueError, match="The optimizer does not seem to reference any XLAFSDP parameter"):
        strategy.setup_optimizer(bad_optimizer)
Test that `setup_optimizer()` validates the param groups and reference to FSDP parameters.
tests/tests_fabric/strategies/test_xla_fsdp.py