Repository: optuna/optuna

Issue:
Index out of bounds error in TpeSampler when multivariate=constant_liar=True and parallel jobs
### Expected behavior

Under a specific configuration using the `TPESampler` (see MWP below), Optuna raises an `IndexError` and stops the study.
I expected Optuna to execute all configured trials without raising an `IndexError`.

### Environment

- Optuna version: 3.1.0
- Python version: 3.8.16
- OS: Linux-5.15.0-60-generic-x86_64-with-glibc2.17
- (Optional) Other libraries and their versions: `conda env export`:
  ```yaml                    
  name: optuna-test
  channels:
    - defaults
  dependencies:
    - _libgcc_mutex=0.1=main
    - _openmp_mutex=5.1=1_gnu
    - ca-certificates=2023.01.10=h06a4308_0
    - certifi=2022.12.7=py38h06a4308_0
    - ld_impl_linux-64=2.38=h1181459_1
    - libffi=3.4.2=h6a678d5_6
    - libgcc-ng=11.2.0=h1234567_1
    - libgomp=11.2.0=h1234567_1
    - libstdcxx-ng=11.2.0=h1234567_1
    - ncurses=6.4=h6a678d5_0
    - openssl=1.1.1t=h7f8727e_0
    - pip=22.3.1=py38h06a4308_0
    - python=3.8.16=h7a1cb2a_2
    - readline=8.2=h5eee18b_0
    - setuptools=65.6.3=py38h06a4308_0
    - sqlite=3.40.1=h5082296_0
    - tk=8.6.12=h1ccaba5_0
    - wheel=0.38.4=py38h06a4308_0
    - xz=5.2.10=h5eee18b_1
    - zlib=1.2.13=h5eee18b_0
    - pip:
        - alembic==1.9.4
        - cmaes==0.9.1
        - colorlog==6.7.0
        - greenlet==2.0.2
        - importlib-metadata==6.0.0
        - importlib-resources==5.12.0
        - mako==1.2.4
        - markupsafe==2.1.2
        - numpy==1.24.2
        - optuna==3.1.0
        - packaging==23.0
        - pyyaml==6.0
        - sqlalchemy==2.0.4
        - tqdm==4.64.1
        - typing-extensions==4.5.0
        - zipp==3.14
  ```


### Error messages, stack traces, or logs

```shell
scripts/optuna_tpe_multi_mwp.py:23: ExperimentalWarning: JournalStorage is experimental (supported from v3.1.0). The interface can change in the future.
  storage = JournalStorage(
<truncated>/python3.8/site-packages/optuna/samplers/_tpe/sampler.py:282: ExperimentalWarning: ``multivariate`` option is an experimental feature. The interface can change in the future.
  warnings.warn(
<truncated>/python3.8/site-packages/optuna/samplers/_tpe/sampler.py:301: ExperimentalWarning: ``constant_liar`` option is an experimental feature. The interface can change in the future.
  warnings.warn(
[I 2023-02-21 17:58:55,568] A new study created in Journal with name: TPESampler with multivariate and constant_liar (seed=42)
[I 2023-02-21 17:58:55,604] Trial 1 finished with value: 0.30471707975443135 and parameters: {'int_param': 1, 'categorical_param': 'value2'}. Best is trial 0 with value: -1.0399841062404955.
[I 2023-02-21 17:58:55,604] Trial 0 finished with value: -1.0399841062404955 and parameters: {'int_param': 6, 'categorical_param': 'value1'}. Best is trial 0 with value: -1.0399841062404955.
[I 2023-02-21 17:58:55,638] Trial 3 finished with value: 0.7504511958064572 and parameters: {'int_param': 6, 'categorical_param': 'value1'}. Best is trial 0 with value: -1.0399841062404955.
[I 2023-02-21 17:58:55,641] Trial 2 finished with value: 0.9405647163912139 and parameters: {'int_param': 2, 'categorical_param': 'value2'}. Best is trial 0 with value: -1.0399841062404955.
[I 2023-02-21 17:58:55,683] Trial 4 finished with value: -1.9510351886538364 and parameters: {'int_param': 5, 'categorical_param': 'value2'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,696] Trial 5 finished with value: -1.302179506862318 and parameters: {'int_param': 3, 'categorical_param': 'value2'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,726] Trial 6 finished with value: 0.12784040316728537 and parameters: {'int_param': 2, 'categorical_param': 'value1'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,753] Trial 7 finished with value: -0.3162425923435822 and parameters: {'int_param': 10, 'categorical_param': 'value1'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,780] Trial 8 finished with value: -0.016801157504288795 and parameters: {'int_param': 10, 'categorical_param': 'value2'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,808] Trial 9 finished with value: -0.85304392757358 and parameters: {'int_param': 9, 'categorical_param': 'value2'}. Best is trial 4 with value: -1.9510351886538364.
[I 2023-02-21 17:58:55,856] Trial 10 finished with value: 0.8793979748628286 and parameters: {'int_param': 5, 'categorical_param': 'value2'}. Best is trial 4 with value: -1.9510351886538364.
Traceback (most recent call last):
  File "scripts/optuna_tpe_multi_mwp.py", line 41, in <module>
    mwp(42 + i)
  File "scripts/optuna_tpe_multi_mwp.py", line 33, in mwp
    study.optimize(objective, n_trials=100, n_jobs=2)
  File "<truncated>/python3.8/site-packages/optuna/study/study.py", line 425, in optimize
    _optimize(
  File "<truncated>/python3.8/site-packages/optuna/study/_optimize.py", line 103, in _optimize
    f.result()
  File "<truncated>/python3.8/concurrent/futures/_base.py", line 437, in result
    return self.__get_result()
  File "<truncated>/python3.8/concurrent/futures/_base.py", line 389, in __get_result
    raise self._exception
  File "<truncated>/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "<truncated>/python3.8/site-packages/optuna/study/_optimize.py", line 163, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
  File "<truncated>/python3.8/site-packages/optuna/study/_optimize.py", line 191, in _run_trial
    trial = study.ask()
  File "<truncated>/python3.8/site-packages/optuna/study/study.py", line 519, in ask
    trial = trial_module.Trial(self, trial_id)
  File "<truncated>/python3.8/site-packages/optuna/trial/_trial.py", line 56, in __init__
    self._init_relative_params()
  File "<truncated>/python3.8/site-packages/optuna/trial/_trial.py", line 66, in _init_relative_params
    self.relative_params = self.study.sampler.sample_relative(
  File "<truncated>/python3.8/site-packages/optuna/samplers/_tpe/sampler.py", line 377, in sample_relative
    return self._sample_relative(study, trial, search_space)
  File "<truncated>/python3.8/site-packages/optuna/samplers/_tpe/sampler.py", line 421, in _sample_relative
    mpe_above = _ParzenEstimator(above, search_space, self._parzen_estimator_parameters)
  File "<truncated>/python3.8/site-packages/optuna/samplers/_tpe/parzen_estimator.py", line 82, in __init__
    categorical_weights = self._calculate_categorical_params(
  File "<truncated>/python3.8/site-packages/optuna/samplers/_tpe/parzen_estimator.py", line 379, in _calculate_categorical_params
    weights[np.arange(n_observations), observations] += 1
IndexError: index -9223372036854775808 is out of bounds for axis 1 with size 2
```


### Steps to reproduce

1. Install optuna 3.1.0
2. Create an objective with at least one categorical and one integer parameter; the result is not relevant
3. Use either the new `JournalFileStorage` or the `RDBStorage`. The `InMemoryStorage` does not raise the error.
4. Create a study with the `TpeSampler` and `multivariate=True, constant_liar=True`
5. Start the optimization in parallel using `n_jobs>=2`.

MWP-code:

```python
from pathlib import Path

import numpy as np
import optuna
from optuna import Trial
from optuna.samplers import TPESampler
from optuna.storages import JournalStorage, JournalFileStorage, JournalFileOpenLock


def mwp(seed: int) -> None:
    rng = np.random.default_rng(seed)

    def objective(trial: Trial) -> float:
        # we need at least one categorical and one integer parameter:
        int_param = trial.suggest_int("int_param", 1, 10)
        categorical_param = trial.suggest_categorical("categorical_param", ["value1", "value2"])

        # result is not relevant:
        return rng.normal()

    # we need a storage that serializes results; the in-memory-storage does not produce the error:
    journal_file_path = Path.cwd() / "test-optuna-journal.log"
    storage = JournalStorage(
        JournalFileStorage(str(journal_file_path), lock_obj=JournalFileOpenLock(str(journal_file_path)))
    )
    study = optuna.create_study(
        study_name=f"TPESampler with multivariate and constant_liar (seed={seed})",
        storage=storage,
        # the multivariate and constant_liar parameters must be set to True:
        sampler=TPESampler(multivariate=True, constant_liar=True, warn_independent_sampling=False),
    )
    # optimize function in parallel (n_jobs >= 2):
    study.optimize(objective, n_trials=100, n_jobs=2)

    print(f"... done! Best params with value={study.best_value}: {study.best_params}")
    journal_file_path.unlink(missing_ok=True)


if __name__ == '__main__':
    mwp(42)
```


### Additional context (optional)

I started to look into the error, but couldn't find the source of it yet. This is the current state of my investigation:

In the following snippet:

https://github.com/optuna/optuna/blob/6adfe8cd953b3b8d85aad3a6c0b0991cbdc6ad40/optuna/samplers/_tpe/sampler.py#L399-L410

- `above` still contains `np.nan`s (last entry in array) for the categorical parameter (`categorical_param`)
- when recalculating the weights later on, this `np.nan` is cast to `int` resulting in an underflow (value `-9223372036854775808`)

But I couldn't figure out why this `np.nan` is still in `above["categorical_param"]` and it's just happening when I use parallelism and a non-in-memory-storage.


Discussion:
I just tested with the [current master branch](https://github.com/optuna/optuna/tree/dd37119abbc7da0420bb9666bb84baed519d4d72) and the issue persists:

```bash
$ python3 -c 'import optuna; print(f"- Optuna version:{optuna.__version__}")'
- Optuna version:3.2.0.dev
$ python3 -c 'import platform; print(f"- Python version:{platform.python_version()}")'
- Python version:3.8.16
$ python3 -c 'import platform; print(f"- OS:{platform.platform()}")'
- OS:Linux-5.15.0-60-generic-x86_64-with-glibc2.17
$ git log -n 1 --oneline
dd37119ab Merge pull request #4442 from g-votte/rename-a-evaluator-arg
```
We appreciate your detailed report and effort to find out the cause.
~~First of all, setting both `multivariate=True` and `constant_liar=True` are incompatible as mentioned in #3565 and [document](https://optuna.readthedocs.io/en/stable/reference/samplers/generated/optuna.samplers.TPESampler.html#:~:text=workers%20is%20high.-,Note,-This%20feature%20can). Also, you have to give `direction` as an input when you run `create_study()` for multi-objective optimizations.~~

~~The following is a possible cause of the error that I found out thanks to the information you provided.
Trials whose state is `TrialState.RUNNING` or `TrialState.PRUNED` is ignored when you use multivariate TPE as you can see below:~~
https://github.com/optuna/optuna/blob/6adfe8cd953b3b8d85aad3a6c0b0991cbdc6ad40/optuna/samplers/_tpe/sampler.py#L641-L658
~~However, `study._is_multi_objective()` is determined by the length of `study.directions` and this returns `False` unless you give `directions` when you create a study. Therefore, unexpected parts seem to have been executed in your case.~~
~~In addition, I am not sure why but `trial.intermediate_values` is not properly propagated between different processes due to the [_CachedStorage](https://github.com/optuna/optuna/blob/dd37119abbc7da0420bb9666bb84baed519d4d72/optuna/storages/_cached_storage.py#L37) (this is a bug). Therefore, I guess these are the reasons why `above` contains `np.nan`. Giving `direction` when you run `create_study()` (and setting `constant_liar=False`) would solve the problem.~~
~~If it is not the case, trials in state `TrialState.RUNNING` might be the cause of the problem because `TPESampler`s only consider trials whose states are `TrialState.RUNNING` when you set `constant_liar=True`(, which is not supported in multivariate TPE anyway).~~
Thank you for your bug report. I found these lines are buggy.

https://github.com/optuna/optuna/blob/1ef80b1913047b83ba169185365b3e1524f920a0/optuna/samplers/_tpe/sampler.py#L393-L401

Minimal reproducible example is here. It's a severe bug, so I will address this problem with high priority.

```python
import time
import optuna

def objective(trial):
    trial.suggest_int("x", 1, 10)
    time.sleep(0.1)
    trial.suggest_categorical("y", ["a", "b"])
    return 0

sampler = optuna.samplers.TPESampler(multivariate=True, constant_liar=True, seed=42)
study = optuna.create_study(sampler=sampler)
study.optimize(objective, n_jobs=2)
```

----

@Alnusjaponica Have you misread multivariate as multiobjective?
Thank you for your indication. You're right and I misunderstood the situation. I apologize for the confusion.

Resolving patch:
diff --git a/optuna/samplers/_tpe/sampler.py b/optuna/samplers/_tpe/sampler.py
--- a/optuna/samplers/_tpe/sampler.py
+++ b/optuna/samplers/_tpe/sampler.py
@@ -392,10 +392,9 @@ def _sample_relative(
         indices_below, indices_above = _split_observation_pairs(scores, self._gamma(n), violations)
         # `None` items are intentionally converted to `nan` and then filtered out.
         # For `nan` conversion, the dtype must be float.
-        # `None` items appear only when `group=True`. We just use the first parameter because the
-        # masks are the same for all parameters in one group.
+        # `None` items appear when `group=True` or `constant_liar=True`.
         config_values = {k: np.asarray(v, dtype=float) for k, v in values.items()}
-        param_mask = ~np.isnan(list(config_values.values())[0])
+        param_mask = np.all(~np.isnan(list(config_values.values())), axis=0)
         param_mask_below, param_mask_above = param_mask[indices_below], param_mask[indices_above]
         below = {k: v[indices_below[param_mask_below]] for k, v in config_values.items()}
         above = {k: v[indices_above[param_mask_above]] for k, v in config_values.items()}


Test patch:
diff --git a/tests/samplers_tests/tpe_tests/test_sampler.py b/tests/samplers_tests/tpe_tests/test_sampler.py
--- a/tests/samplers_tests/tpe_tests/test_sampler.py
+++ b/tests/samplers_tests/tpe_tests/test_sampler.py
@@ -1154,3 +1154,34 @@ def test_constant_liar_observation_pairs(direction: str) -> None:
 def test_constant_liar_experimental_warning() -> None:
     with pytest.warns(optuna.exceptions.ExperimentalWarning):
         _ = TPESampler(constant_liar=True)
+
+
+@pytest.mark.parametrize("multivariate", [True, False])
+def test_constant_liar_with_running_trial(multivariate: bool) -> None:
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", optuna.exceptions.ExperimentalWarning)
+        sampler = TPESampler(multivariate=multivariate, constant_liar=True, n_startup_trials=0)
+
+    study = optuna.create_study(sampler=sampler)
+
+    # Add a complete trial.
+    trial0 = study.ask()
+    trial0.suggest_int("x", 0, 10)
+    trial0.suggest_float("y", 0, 10)
+    trial0.suggest_categorical("z", [0, 1, 2])
+    study.tell(trial0, 0)
+
+    # Add running trials.
+    trial1 = study.ask()
+    trial1.suggest_int("x", 0, 10)
+    trial2 = study.ask()
+    trial2.suggest_float("y", 0, 10)
+    trial3 = study.ask()
+    trial3.suggest_categorical("z", [0, 1, 2])
+
+    # Test suggestion with running trials.
+    trial = study.ask()
+    trial.suggest_int("x", 0, 10)
+    trial.suggest_float("y", 0, 10)
+    trial.suggest_categorical("z", [0, 1, 2])
+    study.tell(trial, 0)