Repository: optuna/optuna

Issue:
Parameter suggestion with PartialFixedSampler and TPE fails mid-study
### Expected behavior

Study runs without errors.

### Environment

- Optuna version:4.7.0
- Python version:3.12.12
- OS:Linux-4.18.0-553.75.1.el8_10.x86_64-x86_64-with-glibc2.39


### Error messages, stack traces, or logs

```shell
Eventually, after several trials have finished, get:


File /opt/conda/envs/XXX/lib/python3.12/site-packages/optuna/trial/_trial.py:163, in Trial.suggest_float(self, name, low, high, step, log)
    162 distribution = FloatDistribution(low, high, log=log, step=step)
--> 163 suggested_value = self._suggest(name, distribution)
    164 self._check_distribution(name, distribution)

File /opt/conda/envs/XXX/lib/python3.12/site-packages/optuna/trial/_trial.py:633, in Trial._suggest(self, name, distribution)
    631 elif distribution.single():
    632     param_value = distributions._get_single_value(distribution)
--> 633 elif self._is_relative_param(name, distribution):
    634     param_value = self.relative_params[name]
    635 else:

File /opt/conda/envs/XXX/lib/python3.12/site-packages/optuna/trial/_trial.py:672, in Trial._is_relative_param(self, name, distribution)
    669 assert self.relative_search_space is not None
    671 if name not in self.relative_search_space:
--> 672     raise ValueError(...)
    677 relative_distribution = self.relative_search_space[name]

ValueError: The parameter 'lr' was sampled by `sample_relative` method but it is not contained in the relative search space.
```

### Steps to reproduce

Unfortunately, so far I could not track down what ultimately triggers this error.
The first few trials (more than `n_startup_trials`) run complete nicely, some are pruned.
The last trial before the error is unremarkable.
At some point, a trial will fail at the first parameter suggestion and all following trials also.

1. Create the study with no trials, save as sqlite db.
2. On a slurm cluster, each job accesses the db, `load_study()` with the same sampler and runs `study.optimize()`
3. The study is loaded with a  `TPESampler(n_startup_trials=20, multivariate=True, group=True)`, which is then replaced by a `PartialFixedSampler` that wraps the `TPESampler`.
4. The parameter in question (`lr`) is one of these fixed parameters. I never do fancy stuff with it, only `trial.suggest_float('lr', 1e-5, 1e-2, log=True)` in each trial.
5. I do fancy stuff with other parameters (conditional on values of other parameters, changing bounds depending on other parameters).

What kinda bugs me is that when I look at `study.get_trials()`, the `distributions` contain the bounds as in the suggest, not the fixed value. Also, from the stack trace, it appears that `lr` is not in `trial._fixed_params`. I wonder if the `PartialFixedSampler` is working as intentional?

I would be happy for any guidance towards narrowing down this issue. Thanks!


### Additional context (optional)

_No response_

Discussion hints:

Solution patch:
diff --git a/optuna/samplers/_tpe/sampler.py b/optuna/samplers/_tpe/sampler.py
index 28d4cd839a..be36eb13c3 100644
--- a/optuna/samplers/_tpe/sampler.py
+++ b/optuna/samplers/_tpe/sampler.py
@@ -418,12 +418,17 @@ def sample_relative(
             assert self._search_space_group is not None
             params = {}
             for sub_space in self._search_space_group.search_spaces:
-                search_space = {}
+                _search_space = {}
                 # Sort keys because Python's string hashing is nondeterministic.
                 for name, distribution in sorted(sub_space.items()):
-                    if not distribution.single():
-                        search_space[name] = distribution
-                params.update(self._sample_relative(study, trial, search_space))
+                    if distribution.single():
+                        continue
+                    if name not in search_space:
+                        # When used together with PartialFixedSampler, the search space may be
+                        # smaller than what is inferred from the study.
+                        continue
+                    _search_space[name] = distribution
+                params.update(self._sample_relative(study, trial, _search_space))
         else:
             params = self._sample_relative(study, trial, search_space)
 


Test patch:
diff --git a/tests/samplers_tests/test_partial_fixed.py b/tests/samplers_tests/test_partial_fixed.py
index d454328d72..02299e5a1b 100644
--- a/tests/samplers_tests/test_partial_fixed.py
+++ b/tests/samplers_tests/test_partial_fixed.py
@@ -1,38 +1,55 @@
+from collections.abc import Callable
 from unittest.mock import patch
 import warnings
 
 import pytest
 
 import optuna
+from optuna.samplers import BaseSampler
 from optuna.samplers import PartialFixedSampler
 from optuna.samplers import RandomSampler
 from optuna.trial import Trial
 
 
-def test_fixed_sampling() -> None:
+parametrize_sampler = pytest.mark.parametrize(
+    "sampler_class",
+    [
+        optuna.samplers.RandomSampler,
+        lambda: optuna.samplers.TPESampler(n_startup_trials=0),
+        lambda: optuna.samplers.TPESampler(n_startup_trials=0, multivariate=True),
+        lambda: optuna.samplers.TPESampler(n_startup_trials=0, multivariate=True, group=True),
+        lambda: optuna.samplers.CmaEsSampler(n_startup_trials=0),
+        lambda: optuna.samplers.CmaEsSampler(n_startup_trials=0, use_separable_cma=True),
+        optuna.samplers.NSGAIISampler,
+        optuna.samplers.NSGAIIISampler,
+        optuna.samplers.QMCSampler,
+        lambda: optuna.samplers.GPSampler(n_startup_trials=0),
+        lambda: optuna.samplers.GPSampler(n_startup_trials=0, deterministic_objective=True),
+    ],
+)
+
+
+@parametrize_sampler
+def test_fixed_sampling(sampler_class: Callable[[], BaseSampler]) -> None:
     def objective(trial: Trial) -> float:
         x = trial.suggest_float("x", -10, 10)
         y = trial.suggest_float("y", -10, 10)
-        return x**2 + y**2
+        z = trial.suggest_float("z", -10, 10)
+        return x**2 + y**2 + z**2
+
+    base_sampler = sampler_class()
 
-    study0 = optuna.create_study()
-    study0.sampler = RandomSampler(seed=42)
-    study0.optimize(objective, n_trials=1)
-    x_sampled0 = study0.trials[0].params["x"]
+    study = optuna.create_study(sampler=base_sampler)
+    study.optimize(objective, n_trials=1)
 
-    # Fix parameter ``y`` as 0.
-    study1 = optuna.create_study()
+    # Fix parameter ``z`` as 0.
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", optuna.exceptions.ExperimentalWarning)
-        study1.sampler = PartialFixedSampler(
-            fixed_params={"y": 0}, base_sampler=RandomSampler(seed=42)
-        )
-    study1.optimize(objective, n_trials=1)
+        sampler = PartialFixedSampler(fixed_params={"z": 0}, base_sampler=base_sampler)
+    study.sampler = sampler
+    study.optimize(objective, n_trials=1)
 
-    x_sampled1 = study1.trials[0].params["x"]
-    y_sampled1 = study1.trials[0].params["y"]
-    assert x_sampled1 == x_sampled0
-    assert y_sampled1 == 0
+    assert study.trials[1].params["z"] == 0
 
 
 def test_float_to_int() -> None: