Repository: optuna/optuna

Issue:
TPE not working with a categorical variable with different choice types
Hi. Thanks for this amazing library!

When I try to optimize a model with a categorical variable with different choice types, it gives me `IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (9,) (4,)`. This error always happens after 10 trials.

I would be very grateful for an explanation on why is that the case - I guess it has to do with how TPE makes categorical variables continuous, but I wasn't able to find how it's done.

Also, I think that it might be very useful for people relatively new to Optuna if the error was more self-explanatory. In my case, it took over 10h to find out why my optimization failed.

## Environment

- Optuna version: 2.10.0
- Python version: 3.6.10
- OS: Linux-4.18.0-305.10.2.el8_4.ppc64le-ppc64le-with-redhat-8.4-Ootpa

## Error messages, stack traces, or logs

```
/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/sampler.py:266: ExperimentalWarning: ``multivariate`` option is an experimental feature. The interface can change in the future.
  ExperimentalWarning,
/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/sampler.py:285: ExperimentalWarning: ``constant_liar`` option is an experimental feature. The interface can change in the future.
  ExperimentalWarning,
[I 2021-11-28 07:31:46,317] Trial 0 finished with value: 51.25318418441549 and parameters: {'x1': 7.159132921270249, 'x2': None}. Best is trial 0 with value: 51.25318418441549.
[I 2021-11-28 07:31:46,661] Trial 1 finished with value: 0.0 and parameters: {'x1': -6.523347451025641, 'x2': 1}. Best is trial 0 with value: 51.25318418441549.
[I 2021-11-28 07:31:46,857] Trial 2 finished with value: 77.13131244750303 and parameters: {'x1': 8.782443421252598, 'x2': None}. Best is trial 2 with value: 77.13131244750303.
[I 2021-11-28 07:31:47,067] Trial 3 finished with value: 71.35511499232025 and parameters: {'x1': -8.447195688056496, 'x2': None}. Best is trial 2 with value: 77.13131244750303.
[I 2021-11-28 07:31:47,283] Trial 4 finished with value: 0.0 and parameters: {'x1': 7.448216057409198, 'x2': 1}. Best is trial 2 with value: 77.13131244750303.
[I 2021-11-28 07:31:47,537] Trial 5 finished with value: 42.52100156846517 and parameters: {'x1': 6.520812953034703, 'x2': None}. Best is trial 2 with value: 77.13131244750303.
[I 2021-11-28 07:31:47,757] Trial 6 finished with value: 0.0 and parameters: {'x1': 3.2600618681215856, 'x2': 1}. Best is trial 2 with value: 77.13131244750303.
[I 2021-11-28 07:31:47,905] Trial 7 finished with value: 88.15753251520141 and parameters: {'x1': 9.389224276541775, 'x2': None}. Best is trial 7 with value: 88.15753251520141.
[I 2021-11-28 07:31:48,029] Trial 8 finished with value: 39.329878182295595 and parameters: {'x1': 6.271353775884087, 'x2': None}. Best is trial 7 with value: 88.15753251520141.
[I 2021-11-28 07:31:48,144] Trial 9 finished with value: 0.0 and parameters: {'x1': -7.077110913961812, 'x2': 1}. Best is trial 7 with value: 88.15753251520141.
Traceback (most recent call last):
  File "optuna_test.py", line 23, in <module>
    study.optimize(train, n_trials=300)
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/study/study.py", line 409, in optimize
    show_progress_bar=show_progress_bar,
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/study/_optimize.py", line 76, in _optimize
    progress_bar=progress_bar,
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/study/_optimize.py", line 163, in _optimize_sequential
    trial = _run_trial(study, func, catch)
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/study/_optimize.py", line 194, in _run_trial
    trial = study.ask()
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/study/study.py", line 486, in ask
    trial = trial_module.Trial(self, trial_id)
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/trial/_trial.py", line 56, in __init__
    self._init_relative_params()
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/trial/_trial.py", line 66, in _init_relative_params
    study, trial, self.relative_search_space
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/sampler.py", line 349, in sample_relative
    return self._sample_relative(study, trial, search_space)
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/sampler.py", line 386, in _sample_relative
    mpe_above = _ParzenEstimator(above, search_space, self._parzen_estimator_parameters)
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/parzen_estimator.py", line 87, in __init__
    param_observations, param_name
  File "/home/chledj01/.local/lib/python3.6/site-packages/optuna/samplers/_tpe/parzen_estimator.py", line 366, in _calculate_categorical_params
    weights[np.arange(n_observations), observations] += 1
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (9,) (4,)
```

## Steps to reproduce

1. Create an optuna study with `optuna create-study --study-name '1234' --storage 'sqlite:///test.db' --direction 'maximize'`
2. Run `optuna_test.py` defined below.

## Reproducible examples (optional)

`optuna_test.py`:
```python
import optuna


def train(trial):
    x1 = trial.suggest_uniform('x1', -10, 10)
    x2 = trial.suggest_categorical('x2', [None, 1])
    if x2:
        score = 0
    else:
        score = x1 ** 2
    return score


if __name__ == "__main__":

    sampler = optuna.samplers.TPESampler(multivariate=True,
                                         group=False,
                                         constant_liar=True)
    study = optuna.load_study(study_name='1234',
                              sampler=sampler,
                              storage="sqlite:///test.db")

    study.optimize(train, n_trials=300)
```


Discussion:
Hello @chledowski, thank you for using Optuna!
I haven't figured out the cause yet but I just noticed the error came with `None` in choices.
(I just updated the line of `suggest_categorical` to be `suggest_categorical("x2", ["None", 1])`)

Does this workaround work for your use case? I suspect it is a bug in `TPESampler` but it has to take some times for further investigation. :bow:
Hi @himkt, thanks for the swift response!

The workaround works well, thanks :)
I figured out these lines are buggy. `raw_param_value is not None` is not equivalent to `param_name in trial.params`. https://github.com/optuna/optuna/blob/828593078770607555a10eccb06d70d68afa05b8/optuna/samplers/_tpe/sampler.py#L632-L640
This issue has not seen any recent activity.

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
@@ -632,9 +632,8 @@ def _get_observation_pairs(
 
         # We extract param_value from the trial.
         for param_name in param_names:
-            raw_param_value = trial.params.get(param_name, None)
             param_value: Optional[float]
-            if raw_param_value is not None:
+            if param_name in trial.params:
                 distribution = trial.distributions[param_name]
                 param_value = distribution.to_internal_repr(trial.params[param_name])
             else:


Test patch: