Repository: tinkoff-ai/etna (language: python)

Issue:
Add distribution classes for AutoML
### 🚀 Feature Request

Now we are using distribution classes from optuna, but all we need from them is their name and parameters. Using external classes complicates their integration into transforms and models.

### Proposal

- Add abstract `BaseDistribution` into `etna.auto.distributions`
- Add `CategoricalDistribution` with `choices` parameter
  - We should validate the types of `choices` like `optuna` does
- Add `IntDistribution` with parameters `low`, `high`, `log`
  - We should validate input parameters like `optuna` does
- Add `FloatDistribution` with parameters `low`, `high`, `log`
  - We should validate input parameters like `optuna` does
- Rewrite `params_to_tune` method to return `Dict[str, BaseDistribution]`
- Fix `Tune` to use new distributions
- Add section about distributions into AutoML notebook

### Test cases

- Fix current tests that work with `optuna` distributions
- Add tests on creation of new tests

### Additional context

_No response_
Add distribution classes for AutoML
### 🚀 Feature Request

Now we are using distribution classes from optuna, but all we need from them is their name and parameters. Using external classes complicates their integration into transforms and models.

### Proposal

- Add abstract `BaseDistribution` into `etna.auto.distributions`
- Add `CategoricalDistribution` with `choices` parameter
  - We should validate the types of `choices` like `optuna` does
- Add `IntDistribution` with parameters `low`, `high`, `log`
  - We should validate input parameters like `optuna` does
- Add `FloatDistribution` with parameters `low`, `high`, `log`
  - We should validate input parameters like `optuna` does
- Rewrite `params_to_tune` method to return `Dict[str, BaseDistribution]`
- Fix `Tune` to use new distributions
- Add section about distributions into AutoML notebook

### Test cases

- Fix current tests that work with `optuna` distributions
- Add tests on creation of new tests

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 88b0f9947..6606644ab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,7 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Parameters `estimate_n_folds` and `context_size` to forecast and backtest CLI commands ([#1284](https://github.com/tinkoff-ai/etna/pull/1284))
 -
 - Class `Tune` for hyperparameter optimization within existing pipeline ([#1200](https://github.com/tinkoff-ai/etna/pull/1200))
--
+- Add `etna.distributions` for using it instead of using `optuna.distributions` ([#1292](https://github.com/tinkoff-ai/etna/pull/1292))
 - 
 ### Changed
 - Set the default value of `final_model` to `LinearRegression(positive=True)` in the constructor of `StackingEnsemble` ([#1238](https://github.com/tinkoff-ai/etna/pull/1238))
diff --git a/docs/source/api.rst b/docs/source/api.rst
index 6b1e20c7e..89f6fafac 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -17,6 +17,7 @@ API
    reconciliation
    analysis
    auto
+   distributions
    clustering
    loggers
    commands
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 6476a639d..a0ade5cd5 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -73,7 +73,7 @@
     "scipy": ("https://docs.scipy.org/doc/scipy/", None),
     "torch": ("https://pytorch.org/docs/stable/", None),
     "pytorch_lightning": ("https://pytorch-lightning.readthedocs.io/en/stable/", None),
-    "optuna": ("https://optuna.readthedocs.io/en/stable/", None)
+    "optuna": ("https://optuna.readthedocs.io/en/v2.10.1/", None)
 }
 
 autodoc_typehints = "both"
diff --git a/etna/auto/auto.py b/etna/auto/auto.py
index 65da9344c..5ea1ee1cf 100644
--- a/etna/auto/auto.py
+++ b/etna/auto/auto.py
@@ -12,13 +12,6 @@
 import optuna
 import pandas as pd
 from hydra_slayer import get_from_params
-from optuna.distributions import BaseDistribution
-from optuna.distributions import CategoricalDistribution
-from optuna.distributions import DiscreteUniformDistribution
-from optuna.distributions import IntLogUniformDistribution
-from optuna.distributions import IntUniformDistribution
-from optuna.distributions import LogUniformDistribution
-from optuna.distributions import UniformDistribution
 from optuna.samplers import BaseSampler
 from optuna.samplers import TPESampler
 from optuna.storages import BaseStorage
@@ -33,7 +26,9 @@
 from etna.auto.runner import AbstractRunner
 from etna.auto.runner import LocalRunner
 from etna.auto.utils import config_hash
+from etna.auto.utils import suggest_parameters
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.loggers import tslogger
 from etna.metrics import MAE
 from etna.metrics import MSE
@@ -760,20 +755,6 @@ def objective(
         objective:
             function that runs specified trial and returns its evaluated score
         """
-        dict_of_distrs = {
-            UniformDistribution: lambda x: ("suggest_uniform", {"low": x.low, "high": x.high}),
-            LogUniformDistribution: lambda x: ("suggest_loguniform", {"low": x.low, "high": x.high}),
-            DiscreteUniformDistribution: lambda x: (
-                "suggest_discrete_uniform",
-                {"low": x.low, "high": x.high, "q": x.q},
-            ),
-            IntUniformDistribution: lambda x: ("suggest_int", {"low": x.low, "high": x.high, "step": x.step}),
-            IntLogUniformDistribution: lambda x: (
-                "suggest_int",
-                {"low": x.low, "high": x.high, "step": x.step},
-            ),
-            CategoricalDistribution: lambda x: ("suggest_categorical", {"choices": x.choices}),
-        }
 
         def _find_duplicate_trial(trial: Trial, pipeline: BasePipeline) -> Optional[FrozenTrial]:
             pipeline_hash = config_hash(pipeline.to_dict())
@@ -788,14 +769,7 @@ def _find_duplicate_trial(trial: Trial, pipeline: BasePipeline) -> Optional[Froz
             return None
 
         def _objective(trial: Trial) -> float:
-            # using received optuna.distribution objects to call corresponding trial.suggest_xxx
-            params_suggested = {}
-            for param_name, param_distr in params_to_tune.items():
-                method_name, method_kwargs = dict_of_distrs[type(param_distr)](param_distr)
-                method = getattr(trial, method_name)
-                params_suggested[param_name] = method(param_name, **method_kwargs)
-
-            # create pipeline instance with the parameters to try
+            params_suggested = suggest_parameters(trial=trial, params_to_tune=params_to_tune)
             pipeline_trial_params: BasePipeline = pipeline.set_params(**params_suggested)
 
             duplicate_trial = _find_duplicate_trial(trial, pipeline_trial_params)
diff --git a/etna/auto/utils.py b/etna/auto/utils.py
index 669fd9043..976e85403 100644
--- a/etna/auto/utils.py
+++ b/etna/auto/utils.py
@@ -1,9 +1,18 @@
 import json
 import random
 import time
+from dataclasses import asdict
 from hashlib import md5
 from typing import Any
 from typing import Callable
+from typing import Dict
+
+import optuna
+
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 
 
 def config_hash(config: dict):
@@ -24,3 +33,21 @@ def retry(func: Callable[..., Any], max_retries: int = 5, sleep_time: int = 10,
                 continue
             else:
                 raise e
+
+
+def suggest_parameters(trial: optuna.Trial, params_to_tune: Dict[str, BaseDistribution]) -> Dict[str, Any]:
+    """Suggest parameters for a trial."""
+    distributions = {
+        CategoricalDistribution: "suggest_categorical",
+        IntDistribution: "suggest_int",
+        FloatDistribution: "suggest_float",
+    }
+
+    params_suggested = {}
+    for param_name, param_distr in params_to_tune.items():
+        method_name = distributions[type(param_distr)]
+        method_kwargs = asdict(param_distr)
+        method = getattr(trial, method_name)
+        params_suggested[param_name] = method(param_name, **method_kwargs)
+
+    return params_suggested
diff --git a/etna/distributions/__init__.py b/etna/distributions/__init__.py
new file mode 100644
index 000000000..b24b0ffea
--- /dev/null
+++ b/etna/distributions/__init__.py
@@ -0,0 +1,4 @@
+from etna.distributions.distributions import BaseDistribution
+from etna.distributions.distributions import CategoricalDistribution
+from etna.distributions.distributions import FloatDistribution
+from etna.distributions.distributions import IntDistribution
diff --git a/etna/distributions/distributions.py b/etna/distributions/distributions.py
new file mode 100644
index 000000000..5bde631f9
--- /dev/null
+++ b/etna/distributions/distributions.py
@@ -0,0 +1,82 @@
+from dataclasses import dataclass
+from typing import Optional
+from typing import Sequence
+from typing import Union
+
+CategoricalChoiceType = Union[None, bool, int, float, str]
+
+
+@dataclass(frozen=True)
+class BaseDistribution:
+    """Base distribution."""
+
+    def __post_init__(self):
+        if self.__class__ == BaseDistribution:
+            raise TypeError("Cannot instantiate abstract class.")
+
+
+@dataclass(frozen=True)
+class CategoricalDistribution(BaseDistribution):
+    """Categorical distribution.
+
+    The input parameters aren't validated.
+
+    Look at :py:meth:`~optuna.trial.Trial.suggest_categorical` to find more about the meaning of parameters.
+
+    Attributes
+    ----------
+    choices:
+        Possible values to take.
+    """
+
+    choices: Sequence[CategoricalChoiceType]
+
+
+@dataclass(frozen=True)
+class IntDistribution(BaseDistribution):
+    """Integer-based distribution.
+
+    The input parameters aren't validated.
+
+    Look at :py:meth:`~optuna.trial.Trial.suggest_int` to find more about the meaning of parameters.
+
+    Attributes
+    ----------
+    low:
+        The lower bound.
+    high:
+        The upper bound.
+    step:
+        The space between possible values.
+    log:
+        The flag of using log domain.
+    """
+
+    low: int
+    high: int
+    step: int = 1
+    log: bool = False
+
+
+@dataclass(frozen=True)
+class FloatDistribution(BaseDistribution):
+    """Float-based distribution.
+
+    Look at :py:meth:`~optuna.trial.Trial.suggest_float` to find more about the meaning of parameters.
+
+    Attributes
+    ----------
+    low:
+        The lower bound.
+    high:
+        The upper bound.
+    step:
+        The space between possible values.
+    log:
+        The flag of using log domain.
+    """
+
+    low: float
+    high: float
+    step: Optional[float] = None
+    log: bool = False
diff --git a/etna/ensembles/direct_ensemble.py b/etna/ensembles/direct_ensemble.py
index 3455091bd..c4469c0ef 100644
--- a/etna/ensembles/direct_ensemble.py
+++ b/etna/ensembles/direct_ensemble.py
@@ -10,15 +10,12 @@
 from joblib import Parallel
 from joblib import delayed
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.ensembles.mixins import EnsembleMixin
 from etna.ensembles.mixins import SaveEnsembleMixin
 from etna.pipeline.base import BasePipeline
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class DirectEnsemble(EnsembleMixin, SaveEnsembleMixin, BasePipeline):
     """DirectEnsemble is a pipeline that forecasts future values merging the forecasts of base pipelines.
@@ -162,7 +159,7 @@ def _predict(
         )
         return prediction
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get hyperparameter grid to tune.
 
         Not implemented for this class.
diff --git a/etna/ensembles/stacking_ensemble.py b/etna/ensembles/stacking_ensemble.py
index 7afc61c7c..82349f6e7 100644
--- a/etna/ensembles/stacking_ensemble.py
+++ b/etna/ensembles/stacking_ensemble.py
@@ -17,17 +17,14 @@
 from sklearn.linear_model import LinearRegression
 from typing_extensions import Literal
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.ensembles.mixins import EnsembleMixin
 from etna.ensembles.mixins import SaveEnsembleMixin
 from etna.loggers import tslogger
 from etna.metrics import MAE
 from etna.pipeline.base import BasePipeline
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class StackingEnsemble(EnsembleMixin, SaveEnsembleMixin, BasePipeline):
     """StackingEnsemble is a pipeline that forecast future using the metamodel to combine the forecasts of the base models.
@@ -268,7 +265,7 @@ def _predict(
         prediction = self._process_forecasts(ts=ts, forecasts=predictions)
         return prediction
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get hyperparameter grid to tune.
 
         Not implemented for this class.
diff --git a/etna/ensembles/voting_ensemble.py b/etna/ensembles/voting_ensemble.py
index 3b3bf72da..6b3ad1f2c 100644
--- a/etna/ensembles/voting_ensemble.py
+++ b/etna/ensembles/voting_ensemble.py
@@ -12,18 +12,15 @@
 from sklearn.ensemble import RandomForestRegressor
 from typing_extensions import Literal
 
-from etna import SETTINGS
 from etna.analysis.feature_relevance.relevance_table import TreeBasedRegressor
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.ensembles.mixins import EnsembleMixin
 from etna.ensembles.mixins import SaveEnsembleMixin
 from etna.loggers import tslogger
 from etna.metrics import MAE
 from etna.pipeline.base import BasePipeline
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class VotingEnsemble(EnsembleMixin, SaveEnsembleMixin, BasePipeline):
     """VotingEnsemble is a pipeline that forecast future values with weighted averaging of it's pipelines forecasts.
@@ -239,7 +236,7 @@ def _predict(
         predictions = self._vote(forecasts=predictions)
         return predictions
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get hyperparameter grid to tune.
 
         Not implemented for this class.
diff --git a/etna/models/base.py b/etna/models/base.py
index c25b79c52..41bc083a1 100644
--- a/etna/models/base.py
+++ b/etna/models/base.py
@@ -17,6 +17,7 @@
 from etna.core import SaveMixin
 from etna.core.mixins import BaseMixin
 from etna.datasets.tsdataset import TSDataset
+from etna.distributions import BaseDistribution
 from etna.loggers import tslogger
 from etna.models.decorators import log_decorator
 from etna.models.mixins import SaveNNMixin
@@ -34,9 +35,6 @@
     LightningModule = Mock  # type: ignore
     SaveNNMixin = Mock  # type: ignore
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class AbstractModel(SaveMixin, ABC, BaseMixin):
     """Interface for model with fit method."""
@@ -82,7 +80,7 @@ def get_model(self) -> Union[Any, Dict[str, Any]]:
         """
         pass
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get grid for tuning hyperparameters.
 
         This is default implementation with empty grid.
diff --git a/etna/models/catboost.py b/etna/models/catboost.py
index d7e6e1263..625ec9d90 100644
--- a/etna/models/catboost.py
+++ b/etna/models/catboost.py
@@ -7,18 +7,15 @@
 from catboost import CatBoostRegressor
 from catboost import Pool
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import BaseAdapter
 from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
 from etna.models.mixins import MultiSegmentModelMixin
 from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin
 from etna.models.mixins import PerSegmentModelMixin
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import LogUniformDistribution
-
 
 class _CatBoostAdapter(BaseAdapter):
     def __init__(
@@ -159,7 +156,7 @@ def predict_components(self, df: pd.DataFrame) -> pd.DataFrame:
 
         return pd.DataFrame(data=components, columns=component_names)
 
-    def _params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def _params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``learning_rate``, ``depth``, ``random_strength``, ``l2_leaf_reg``.
@@ -171,10 +168,10 @@ def _params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "learning_rate": LogUniformDistribution(low=1e-4, high=0.5),
-            "depth": IntUniformDistribution(low=1, high=11, step=1),
-            "random_strength": LogUniformDistribution(low=1e-5, high=10),
-            "l2_leaf_reg": LogUniformDistribution(low=0.1, high=200),
+            "learning_rate": FloatDistribution(low=1e-4, high=0.5, log=True),
+            "depth": IntDistribution(low=1, high=11),
+            "random_strength": FloatDistribution(low=1e-5, high=10, log=True),
+            "l2_leaf_reg": FloatDistribution(low=0.1, high=200, log=True),
         }
 
 
@@ -304,7 +301,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``learning_rate``, ``depth``, ``random_strength``, ``l2_leaf_reg``.
@@ -444,7 +441,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``learning_rate``, ``depth``, ``random_strength``, ``l2_leaf_reg``.
diff --git a/etna/models/deadline_ma.py b/etna/models/deadline_ma.py
index 8d2fac2b6..6890587b4 100644
--- a/etna/models/deadline_ma.py
+++ b/etna/models/deadline_ma.py
@@ -8,14 +8,11 @@
 import pandas as pd
 from typing_extensions import assert_never
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class SeasonalityMode(str, Enum):
     """Enum for seasonality mode for DeadlineMovingAverageModel."""
@@ -391,7 +388,7 @@ def predict(self, ts: TSDataset, prediction_size: int, return_components: bool =
 
         return ts
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``window`` parameter. Other parameters are expected to be set by the user.
@@ -401,7 +398,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         :
             Grid to tune.
         """
-        return {"window": IntUniformDistribution(low=1, high=10, step=1)}
+        return {"window": IntDistribution(low=1, high=10)}
 
 
 __all__ = ["DeadlineMovingAverageModel"]
diff --git a/etna/models/holt_winters.py b/etna/models/holt_winters.py
index 52d62e1bf..6f9d7aeb5 100644
--- a/etna/models/holt_winters.py
+++ b/etna/models/holt_winters.py
@@ -13,7 +13,8 @@
 from statsmodels.tsa.holtwinters import ExponentialSmoothing
 from statsmodels.tsa.holtwinters.results import HoltWintersResultsWrapper
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.models.base import BaseAdapter
 from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
 from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin
@@ -22,10 +23,6 @@
 from etna.models.utils import determine_num_steps
 from etna.models.utils import select_observations
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class _HoltWintersAdapter(BaseAdapter):
     """
@@ -653,7 +650,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``trend``, ``damped_trend``, ``use_boxcox``.
@@ -794,7 +791,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         Returns
diff --git a/etna/models/linear.py b/etna/models/linear.py
index e16fdfae9..33696a138 100644
--- a/etna/models/linear.py
+++ b/etna/models/linear.py
@@ -4,34 +4,23 @@
 from sklearn.linear_model import ElasticNet
 from sklearn.linear_model import LinearRegression
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
 from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
 from etna.models.mixins import MultiSegmentModelMixin
 from etna.models.mixins import NonPredictionIntervalContextIgnorantModelMixin
 from etna.models.mixins import PerSegmentModelMixin
 from etna.models.sklearn import _SklearnAdapter
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import LogUniformDistribution
-    from optuna.distributions import UniformDistribution
-else:
-    from unittest.mock import Mock
-
-    CategoricalDistribution = Mock  # type: ignore
-    UniformDistribution = Mock  # type: ignore
-    LogUniformDistribution = Mock  # type: ignore
-
-
-_LINEAR_GRID: Dict[str, "BaseDistribution"] = {
+_LINEAR_GRID: Dict[str, BaseDistribution] = {
     "fit_intercept": CategoricalDistribution([False, True]),
 }
 
-_ELASTIC_GRID: Dict[str, "BaseDistribution"] = {
+_ELASTIC_GRID: Dict[str, BaseDistribution] = {
     "fit_intercept": CategoricalDistribution([False, True]),
-    "l1_ratio": UniformDistribution(0, 1),
-    "alpha": LogUniformDistribution(low=1e-5, high=1e3),
+    "l1_ratio": FloatDistribution(low=0, high=1),
+    "alpha": FloatDistribution(low=1e-5, high=1e3, log=True),
 }
 
 
@@ -90,7 +79,7 @@ def __init__(self, fit_intercept: bool = True, **kwargs):
             base_model=_LinearAdapter(regressor=LinearRegression(fit_intercept=self.fit_intercept, **self.kwargs))
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         Returns
@@ -153,7 +142,7 @@ def __init__(self, alpha: float = 1.0, l1_ratio: float = 0.5, fit_intercept: boo
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         Returns
@@ -193,7 +182,7 @@ def __init__(self, fit_intercept: bool = True, **kwargs):
             base_model=_LinearAdapter(regressor=LinearRegression(fit_intercept=self.fit_intercept, **self.kwargs))
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         Returns
@@ -256,7 +245,7 @@ def __init__(self, alpha: float = 1.0, l1_ratio: float = 0.5, fit_intercept: boo
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         Returns
diff --git a/etna/models/naive.py b/etna/models/naive.py
index c310519d3..9d0fdc473 100644
--- a/etna/models/naive.py
+++ b/etna/models/naive.py
@@ -1,11 +1,8 @@
 from typing import Dict
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
 from etna.models.seasonal_ma import SeasonalMovingAverageModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class NaiveModel(SeasonalMovingAverageModel):
     """Naive model predicts t-th value of series with its (t - lag) value.
@@ -33,7 +30,7 @@ def __init__(self, lag: int = 1):
         self.lag = lag
         super().__init__(window=1, seasonality=lag)
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid is empty.
diff --git a/etna/models/nn/deepar.py b/etna/models/nn/deepar.py
index 40a69ef6b..91c08bb7e 100644
--- a/etna/models/nn/deepar.py
+++ b/etna/models/nn/deepar.py
@@ -8,6 +8,9 @@
 
 from etna import SETTINGS
 from etna.datasets.tsdataset import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import PredictionIntervalContextRequiredAbstractModel
 from etna.models.base import log_decorator
 from etna.models.mixins import SaveNNMixin
@@ -23,12 +26,6 @@
     from pytorch_forecasting.models import DeepAR
     from pytorch_lightning import LightningModule
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import LogUniformDistribution
-    from optuna.distributions import UniformDistribution
-
 
 class DeepARModel(_DeepCopyMixin, PytorchForecastingMixin, SaveNNMixin, PredictionIntervalContextRequiredAbstractModel):
     """Wrapper for :py:class:`pytorch_forecasting.models.deepar.DeepAR`.
@@ -255,7 +252,7 @@ def get_model(self) -> Any:
         """
         return self.model
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``hidden_size``, ``rnn_layers``, ``dropout``, ``lr``.
@@ -267,8 +264,8 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "hidden_size": IntUniformDistribution(low=4, high=64, step=4),
-            "rnn_layers": IntUniformDistribution(low=1, high=3, step=1),
-            "dropout": UniformDistribution(low=0, high=0.5),
-            "lr": LogUniformDistribution(low=1e-5, high=1e-2),
+            "hidden_size": IntDistribution(low=4, high=64, step=4),
+            "rnn_layers": IntDistribution(low=1, high=3),
+            "dropout": FloatDistribution(low=0, high=0.5),
+            "lr": FloatDistribution(low=1e-5, high=1e-2, log=True),
         }
diff --git a/etna/models/nn/mlp.py b/etna/models/nn/mlp.py
index e66ecf0eb..81c0034e7 100644
--- a/etna/models/nn/mlp.py
+++ b/etna/models/nn/mlp.py
@@ -9,6 +9,9 @@
 from typing_extensions import TypedDict
 
 from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import DeepBaseModel
 from etna.models.base import DeepBaseNet
 
@@ -16,11 +19,6 @@
     import torch
     import torch.nn as nn
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import LogUniformDistribution
-
 
 class MLPBatch(TypedDict):
     """Batch specification for MLP."""
@@ -245,7 +243,7 @@ def __init__(
             split_params=split_params,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``lr``, ``hidden_size.i`` where i from 0 to ``len(hidden_size) - 1``.
@@ -260,8 +258,8 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
 
         for i in range(len(self.hidden_size)):
             key = f"hidden_size.{i}"
-            value = IntUniformDistribution(low=4, high=64, step=4)
+            value = IntDistribution(low=4, high=64, step=4)
             grid[key] = value
 
-        grid["lr"] = LogUniformDistribution(low=1e-5, high=1e-2)
+        grid["lr"] = FloatDistribution(low=1e-5, high=1e-2, log=True)
         return grid
diff --git a/etna/models/nn/rnn.py b/etna/models/nn/rnn.py
index d544a0b76..1d7c7be09 100644
--- a/etna/models/nn/rnn.py
+++ b/etna/models/nn/rnn.py
@@ -8,6 +8,9 @@
 from typing_extensions import TypedDict
 
 from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import DeepBaseModel
 from etna.models.base import DeepBaseNet
 
@@ -15,11 +18,6 @@
     import torch
     import torch.nn as nn
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import LogUniformDistribution
-
 
 class RNNBatch(TypedDict):
     """Batch specification for RNN."""
@@ -290,7 +288,7 @@ def __init__(
             split_params=split_params,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``num_layers``, ``hidden_size``, ``lr``, ``encoder_length``.
@@ -302,8 +300,8 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "num_layers": IntUniformDistribution(low=1, high=3),
-            "hidden_size": IntUniformDistribution(low=4, high=64, step=4),
-            "lr": LogUniformDistribution(low=1e-5, high=1e-2),
-            "encoder_length": IntUniformDistribution(low=1, high=20),
+            "num_layers": IntDistribution(low=1, high=3),
+            "hidden_size": IntDistribution(low=4, high=64, step=4),
+            "lr": FloatDistribution(low=1e-5, high=1e-2, log=True),
+            "encoder_length": IntDistribution(low=1, high=20),
         }
diff --git a/etna/models/nn/tft.py b/etna/models/nn/tft.py
index d2c9d09f4..006daddcb 100644
--- a/etna/models/nn/tft.py
+++ b/etna/models/nn/tft.py
@@ -9,6 +9,9 @@
 
 from etna import SETTINGS
 from etna.datasets.tsdataset import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import PredictionIntervalContextRequiredAbstractModel
 from etna.models.base import log_decorator
 from etna.models.mixins import SaveNNMixin
@@ -23,12 +26,6 @@
     from pytorch_forecasting.models import TemporalFusionTransformer
     from pytorch_lightning import LightningModule
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import LogUniformDistribution
-    from optuna.distributions import UniformDistribution
-
 
 class TFTModel(_DeepCopyMixin, PytorchForecastingMixin, SaveNNMixin, PredictionIntervalContextRequiredAbstractModel):
     """Wrapper for :py:class:`pytorch_forecasting.models.temporal_fusion_transformer.TemporalFusionTransformer`.
@@ -284,7 +281,7 @@ def get_model(self) -> Any:
         """
         return self.model
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``hidden_size``, ``lstm_layers``, ``dropout``, ``attention_head_size``, ``lr``.
@@ -296,9 +293,9 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "hidden_size": IntUniformDistribution(low=4, high=64, step=4),
-            "lstm_layers": IntUniformDistribution(low=1, high=3, step=1),
-            "dropout": UniformDistribution(low=0, high=0.5),
-            "attention_head_size": IntUniformDistribution(low=2, high=8, step=2),
-            "lr": LogUniformDistribution(low=1e-5, high=1e-2),
+            "hidden_size": IntDistribution(low=4, high=64, step=4),
+            "lstm_layers": IntDistribution(low=1, high=3),
+            "dropout": FloatDistribution(low=0, high=0.5),
+            "attention_head_size": IntDistribution(low=2, high=8, step=2),
+            "lr": FloatDistribution(low=1e-5, high=1e-2, log=True),
         }
diff --git a/etna/models/prophet.py b/etna/models/prophet.py
index 1e30aaede..0e2d8a3fd 100644
--- a/etna/models/prophet.py
+++ b/etna/models/prophet.py
@@ -11,6 +11,9 @@
 import pandas as pd
 
 from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
 from etna.models.base import BaseAdapter
 from etna.models.base import PredictionIntervalContextIgnorantAbstractModel
 from etna.models.mixins import PerSegmentModelMixin
@@ -21,12 +24,6 @@
     from prophet.serialize import model_from_dict
     from prophet.serialize import model_to_dict
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import LogUniformDistribution
-    from optuna.distributions import UniformDistribution
-
 
 class _ProphetAdapter(BaseAdapter):
     """Class for holding Prophet model."""
@@ -458,7 +455,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``seasonality_mode``, ``seasonality_prior_scale``, ``changepoint_prior_scale``,
@@ -472,8 +469,8 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         """
         return {
             "seasonality_mode": CategoricalDistribution(["additive", "multiplicative"]),
-            "seasonality_prior_scale": LogUniformDistribution(low=1e-2, high=10),
-            "changepoint_prior_scale": LogUniformDistribution(low=1e-3, high=0.5),
-            "changepoint_range": UniformDistribution(low=0.8, high=0.95),
-            "holidays_prior_scale": LogUniformDistribution(low=1e-2, high=10),
+            "seasonality_prior_scale": FloatDistribution(low=1e-2, high=10, log=True),
+            "changepoint_prior_scale": FloatDistribution(low=1e-3, high=0.5, log=True),
+            "changepoint_range": FloatDistribution(low=0.8, high=0.95),
+            "holidays_prior_scale": FloatDistribution(low=1e-2, high=10, log=True),
         }
diff --git a/etna/models/sarimax.py b/etna/models/sarimax.py
index fd27fa44d..9f0407cf4 100644
--- a/etna/models/sarimax.py
+++ b/etna/models/sarimax.py
@@ -14,7 +14,9 @@
 from statsmodels.tsa.statespace.sarimax import SARIMAXResultsWrapper
 from statsmodels.tsa.statespace.simulation_smoother import SimulationSmoother
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.libs.pmdarima_utils import seasonal_prediction_with_confidence
 from etna.models.base import BaseAdapter
 from etna.models.base import PredictionIntervalContextIgnorantAbstractModel
@@ -24,11 +26,6 @@
 from etna.models.utils import determine_num_steps
 from etna.models.utils import select_observations
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 warnings.filterwarnings(
     message="No frequency information was provided, so inferred frequency .* will be used",
     action="ignore",
@@ -732,7 +729,7 @@ def __init__(
             )
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``order.0``, ``order.1``, ``order.2``, ``trend``.
@@ -746,9 +743,9 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         grid: Dict[str, "BaseDistribution"] = {
-            "order.0": IntUniformDistribution(low=1, high=6, step=1),
-            "order.1": IntUniformDistribution(low=1, high=2, step=1),
-            "order.2": IntUniformDistribution(low=1, high=6, step=1),
+            "order.0": IntDistribution(low=1, high=6),
+            "order.1": IntDistribution(low=1, high=2),
+            "order.2": IntDistribution(low=1, high=6),
             "trend": CategoricalDistribution(["n", "c", "t", "ct"]),
         }
 
@@ -756,9 +753,9 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         if num_periods > 0:
             grid.update(
                 {
-                    "seasonal_order.0": IntUniformDistribution(low=0, high=2, step=1),
-                    "seasonal_order.1": IntUniformDistribution(low=0, high=1, step=1),
-                    "seasonal_order.2": IntUniformDistribution(low=0, high=1, step=1),
+                    "seasonal_order.0": IntDistribution(low=0, high=2),
+                    "seasonal_order.1": IntDistribution(low=0, high=1),
+                    "seasonal_order.2": IntDistribution(low=0, high=1),
                 }
             )
 
diff --git a/etna/models/seasonal_ma.py b/etna/models/seasonal_ma.py
index 95858de63..a8277604b 100644
--- a/etna/models/seasonal_ma.py
+++ b/etna/models/seasonal_ma.py
@@ -4,14 +4,11 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import IntDistribution
 from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class SeasonalMovingAverageModel(
     NonPredictionIntervalContextRequiredAbstractModel,
@@ -239,7 +236,7 @@ def predict(self, ts: TSDataset, prediction_size: int, return_components: bool =
             ts.add_target_components(target_components_df=target_components_df)
         return ts
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``window`` parameter. Other parameters are expected to be set by the user.
@@ -249,7 +246,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         :
             Grid to tune.
         """
-        return {"window": IntUniformDistribution(low=1, high=10, step=1)}
+        return {"window": IntDistribution(low=1, high=10)}
 
 
 __all__ = ["SeasonalMovingAverageModel"]
diff --git a/etna/pipeline/base.py b/etna/pipeline/base.py
index b502cd137..a7bf71eb4 100644
--- a/etna/pipeline/base.py
+++ b/etna/pipeline/base.py
@@ -19,18 +19,15 @@
 from typing_extensions import TypedDict
 from typing_extensions import assert_never
 
-from etna import SETTINGS
 from etna.core import AbstractSaveable
 from etna.core import BaseMixin
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.loggers import tslogger
 from etna.metrics import MAE
 from etna.metrics import Metric
 from etna.metrics import MetricAggregationMode
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 Timestamp = Union[str, pd.Timestamp]
 
 
@@ -278,7 +275,7 @@ def backtest(
         """
 
     @abstractmethod
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get hyperparameter grid to tune.
 
         Returns
diff --git a/etna/pipeline/mixins.py b/etna/pipeline/mixins.py
index a117e8be7..fe1edaa69 100644
--- a/etna/pipeline/mixins.py
+++ b/etna/pipeline/mixins.py
@@ -11,10 +11,10 @@
 from typing_extensions import Self
 from typing_extensions import get_args
 
-from etna import SETTINGS
 from etna.core import SaveMixin
 from etna.core import load
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.models import ModelType
 from etna.models import NonPredictionIntervalContextIgnorantAbstractModel
 from etna.models import NonPredictionIntervalContextRequiredAbstractModel
@@ -23,9 +23,6 @@
 from etna.models import PredictionIntervalContextRequiredAbstractModel
 from etna.transforms import Transform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class ModelPipelinePredictMixin:
     """Mixin for pipelines with model inside with implementation of ``_predict`` method."""
@@ -117,7 +114,7 @@ class ModelPipelineParamsToTuneMixin:
     model: ModelType
     transforms: Sequence[Transform]
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get hyperparameter grid to tune.
 
         Parameters for model has prefix "model.", e.g. "model.alpha".
diff --git a/etna/transforms/base.py b/etna/transforms/base.py
index 30dd85e48..7f99222b5 100644
--- a/etna/transforms/base.py
+++ b/etna/transforms/base.py
@@ -10,15 +10,12 @@
 import pandas as pd
 from typing_extensions import Literal
 
-from etna import SETTINGS
 from etna.core import BaseMixin
 from etna.core import SaveMixin
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-
 
 class FutureMixin:
     """Mixin for transforms that can convert non-regressor column to a regressor one."""
@@ -165,7 +162,7 @@ def inverse_transform(self, ts: TSDataset) -> TSDataset:
         """
         pass
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get grid for tuning hyperparameters.
 
         This is default implementation with empty grid.
diff --git a/etna/transforms/decomposition/change_points_based/detrend.py b/etna/transforms/decomposition/change_points_based/detrend.py
index d246b8a10..d68b22585 100644
--- a/etna/transforms/decomposition/change_points_based/detrend.py
+++ b/etna/transforms/decomposition/change_points_based/detrend.py
@@ -6,7 +6,9 @@
 from ruptures.detection import Binseg
 from sklearn.linear_model import LinearRegression
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.decomposition.change_points_based.base import ReversibleChangePointsTransform
 from etna.transforms.decomposition.change_points_based.base import _OneSegmentChangePointsTransform
 from etna.transforms.decomposition.change_points_based.change_points_models import BaseChangePointsModelAdapter
@@ -15,11 +17,6 @@
 from etna.transforms.decomposition.change_points_based.per_interval_models import SklearnRegressionPerIntervalModel
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _OneSegmentChangePointsTrendTransform(_OneSegmentChangePointsTransform):
     """_OneSegmentChangePointsTransform subtracts multiple linear trend from series."""
@@ -110,7 +107,7 @@ def _is_change_points_model_default(self) -> bool:
         # it can't see the difference between Binseg(model="ar") and Binseg(model="l1")
         return self.change_points_model.to_dict() == self._default_change_points_model.to_dict()
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         If ``self.change_points_model`` is equal to default then this grid tunes parameters:
@@ -127,7 +124,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
                 "change_points_model.change_points_model.model": CategoricalDistribution(
                     ["l1", "l2", "normal", "rbf", "cosine", "linear", "clinear", "ar", "mahalanobis", "rank"]
                 ),
-                "change_points_model.n_bkps": IntUniformDistribution(low=5, high=30),
+                "change_points_model.n_bkps": IntDistribution(low=5, high=30),
             }
         else:
             return {}
diff --git a/etna/transforms/decomposition/change_points_based/level.py b/etna/transforms/decomposition/change_points_based/level.py
index 3872f6a4d..a6b6f4a13 100644
--- a/etna/transforms/decomposition/change_points_based/level.py
+++ b/etna/transforms/decomposition/change_points_based/level.py
@@ -3,7 +3,9 @@
 
 from ruptures import Binseg
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.decomposition.change_points_based.base import BaseChangePointsModelAdapter
 from etna.transforms.decomposition.change_points_based.base import ReversibleChangePointsTransform
 from etna.transforms.decomposition.change_points_based.change_points_models.ruptures_based import (
@@ -13,11 +15,6 @@
 from etna.transforms.decomposition.change_points_based.per_interval_models import MeanPerIntervalModel
 from etna.transforms.decomposition.change_points_based.per_interval_models import StatisticsPerIntervalModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _OneSegmentChangePointsLevelTransform(_OneSegmentChangePointsTrendTransform):
     def __init__(
@@ -108,7 +105,7 @@ def _is_change_points_model_default(self) -> bool:
         # it can't see the difference between Binseg(model="ar") and Binseg(model="l1")
         return self.change_points_model.to_dict() == self._default_change_points_model.to_dict()
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         If ``self.change_points_model`` is equal to default then this grid tunes parameters:
@@ -125,7 +122,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
                 "change_points_model.change_points_model.model": CategoricalDistribution(
                     ["l1", "l2", "normal", "rbf", "cosine", "linear", "clinear", "ar", "mahalanobis", "rank"]
                 ),
-                "change_points_model.n_bkps": IntUniformDistribution(low=5, high=30),
+                "change_points_model.n_bkps": IntDistribution(low=5, high=30),
             }
         else:
             return {}
diff --git a/etna/transforms/decomposition/change_points_based/segmentation.py b/etna/transforms/decomposition/change_points_based/segmentation.py
index 93aa366f4..812fd8d97 100644
--- a/etna/transforms/decomposition/change_points_based/segmentation.py
+++ b/etna/transforms/decomposition/change_points_based/segmentation.py
@@ -5,7 +5,9 @@
 import pandas as pd
 from ruptures import Binseg
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.decomposition.change_points_based.base import IrreversibleChangePointsTransform
 from etna.transforms.decomposition.change_points_based.base import _OneSegmentChangePointsTransform
 from etna.transforms.decomposition.change_points_based.change_points_models import BaseChangePointsModelAdapter
@@ -14,11 +16,6 @@
 )
 from etna.transforms.decomposition.change_points_based.per_interval_models import ConstantPerIntervalModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _OneSegmentChangePointsSegmentationTransform(_OneSegmentChangePointsTransform):
     """_OneSegmentChangePointsSegmentationTransform make label encoder to change points."""
@@ -113,7 +110,7 @@ def _is_change_points_model_default(self) -> bool:
         # it can't see the difference between Binseg(model="ar") and Binseg(model="l1")
         return self.change_points_model.to_dict() == self._default_change_points_model.to_dict()
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         If ``self.change_points_model`` is equal to default then this grid tunes parameters:
@@ -130,7 +127,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
                 "change_points_model.change_points_model.model": CategoricalDistribution(
                     ["l1", "l2", "normal", "rbf", "cosine", "linear", "clinear", "ar", "mahalanobis", "rank"]
                 ),
-                "change_points_model.n_bkps": IntUniformDistribution(low=5, high=30),
+                "change_points_model.n_bkps": IntDistribution(low=5, high=30),
             }
         else:
             return {}
diff --git a/etna/transforms/decomposition/change_points_based/trend.py b/etna/transforms/decomposition/change_points_based/trend.py
index 87fb0da9c..693bce140 100644
--- a/etna/transforms/decomposition/change_points_based/trend.py
+++ b/etna/transforms/decomposition/change_points_based/trend.py
@@ -5,7 +5,9 @@
 from ruptures import Binseg
 from sklearn.linear_model import LinearRegression
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.decomposition.change_points_based.base import IrreversibleChangePointsTransform
 from etna.transforms.decomposition.change_points_based.change_points_models import BaseChangePointsModelAdapter
 from etna.transforms.decomposition.change_points_based.change_points_models.ruptures_based import (
@@ -15,11 +17,6 @@
 from etna.transforms.decomposition.change_points_based.per_interval_models import PerIntervalModel
 from etna.transforms.decomposition.change_points_based.per_interval_models import SklearnRegressionPerIntervalModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _OneSegmentTrendTransform(_OneSegmentChangePointsTrendTransform):
     """_OneSegmentTrendTransform adds trend as a feature."""
@@ -128,7 +125,7 @@ def _is_change_points_model_default(self) -> bool:
         # it can't see the difference between Binseg(model="ar") and Binseg(model="l1")
         return self.change_points_model.to_dict() == self._default_change_points_model.to_dict()
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         If ``self.change_points_model`` is equal to default then this grid tunes parameters:
@@ -145,7 +142,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
                 "change_points_model.change_points_model.model": CategoricalDistribution(
                     ["l1", "l2", "normal", "rbf", "cosine", "linear", "clinear", "ar", "mahalanobis", "rank"]
                 ),
-                "change_points_model.n_bkps": IntUniformDistribution(low=5, high=30),
+                "change_points_model.n_bkps": IntDistribution(low=5, high=30),
             }
         else:
             return {}
diff --git a/etna/transforms/decomposition/detrend.py b/etna/transforms/decomposition/detrend.py
index 8082154e6..5179a93b7 100644
--- a/etna/transforms/decomposition/detrend.py
+++ b/etna/transforms/decomposition/detrend.py
@@ -9,15 +9,12 @@
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import PolynomialFeatures
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.base import OneSegmentTransform
 from etna.transforms.base import ReversiblePerSegmentWrapper
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _OneSegmentLinearTrendBaseTransform(OneSegmentTransform):
     """Transform for one segment that implements trend subtraction and reconstruction feature."""
@@ -193,7 +190,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "poly_degree": IntUniformDistribution(low=1, high=2),
+            "poly_degree": IntDistribution(low=1, high=2),
         }
 
 
@@ -253,5 +250,5 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "poly_degree": IntUniformDistribution(low=1, high=2),
+            "poly_degree": IntDistribution(low=1, high=2),
         }
diff --git a/etna/transforms/decomposition/stl.py b/etna/transforms/decomposition/stl.py
index 554bb14b6..4ff997d1a 100644
--- a/etna/transforms/decomposition/stl.py
+++ b/etna/transforms/decomposition/stl.py
@@ -11,15 +11,12 @@
 from statsmodels.tsa.forecasting.stl import STLForecast
 from statsmodels.tsa.forecasting.stl import STLForecastResults
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import OneSegmentTransform
 from etna.transforms.base import ReversiblePerSegmentWrapper
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class _OneSegmentSTLTransform(OneSegmentTransform):
     def __init__(
@@ -230,7 +227,7 @@ def get_regressors_info(self) -> List[str]:
         """Return the list with regressors created by the transform."""
         return []
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``model``, ``robust``. Other parameters are expected to be set by the user.
diff --git a/etna/transforms/encoders/categorical.py b/etna/transforms/encoders/categorical.py
index 4706e1d06..135d3ce24 100644
--- a/etna/transforms/encoders/categorical.py
+++ b/etna/transforms/encoders/categorical.py
@@ -9,14 +9,11 @@
 from sklearn.utils._encode import _check_unknown
 from sklearn.utils._encode import _encode
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import IrreversibleTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class ImputerMode(str, Enum):
     """Enum for different imputation strategy."""
@@ -136,7 +133,7 @@ def _get_column_name(self) -> str:
             return self.out_column
         return self.__repr__()
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``strategy`` parameter. Other parameters are expected to be set by the user.
diff --git a/etna/transforms/feature_selection/feature_importance.py b/etna/transforms/feature_selection/feature_importance.py
index 575b3de8e..5d5ae553e 100644
--- a/etna/transforms/feature_selection/feature_importance.py
+++ b/etna/transforms/feature_selection/feature_importance.py
@@ -14,19 +14,15 @@
 from sklearn.tree import ExtraTreeRegressor
 from typing_extensions import Literal
 
-from etna import SETTINGS
 from etna.analysis import RelevanceTable
 from etna.analysis.feature_selection.mrmr_selection import AggregationMode
 from etna.analysis.feature_selection.mrmr_selection import mrmr
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.feature_selection import BaseFeatureSelectionTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
-
 TreeBasedRegressor = Union[
     DecisionTreeRegressor,
     ExtraTreeRegressor,
@@ -142,7 +138,7 @@ def _fit(self, df: pd.DataFrame) -> "TreeFeatureSelectionTransform":
         self.selected_features = self._select_top_k_features(weights, self.top_k)
         return self
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``model``, ``top_k``. Other parameters are expected to be set by the user.
@@ -157,7 +153,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         """
         return {
             "model": CategoricalDistribution(["catboost", "random_forest"]),
-            "top_k": IntUniformDistribution(low=1, high=self.top_k),
+            "top_k": IntDistribution(low=1, high=self.top_k),
         }
 
 
@@ -242,7 +238,7 @@ def _fit(self, df: pd.DataFrame) -> "MRMRFeatureSelectionTransform":
         )
         return self
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``top_k`` parameter. Other parameters are expected to be set by the user.
@@ -255,5 +251,5 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "top_k": IntUniformDistribution(low=1, high=self.top_k),
+            "top_k": IntDistribution(low=1, high=self.top_k),
         }
diff --git a/etna/transforms/feature_selection/gale_shapley.py b/etna/transforms/feature_selection/gale_shapley.py
index b57aa72ab..256f902de 100644
--- a/etna/transforms/feature_selection/gale_shapley.py
+++ b/etna/transforms/feature_selection/gale_shapley.py
@@ -8,16 +8,13 @@
 import pandas as pd
 from typing_extensions import Literal
 
-from etna import SETTINGS
 from etna.analysis import RelevanceTable
 from etna.core import BaseMixin
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.feature_selection.base import BaseFeatureSelectionTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class BaseGaleShapley(BaseMixin):
     """Base class for a member of Gale-Shapley matching."""
@@ -392,7 +389,7 @@ def _fit(self, df: pd.DataFrame) -> "GaleShapleyFeatureSelectionTransform":
             )
         return self
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``top_k``, ``use_rank``. Other parameters are expected to be set by the user.
@@ -405,6 +402,6 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "top_k": IntUniformDistribution(low=1, high=self.top_k),
+            "top_k": IntDistribution(low=1, high=self.top_k),
             "use_rank": CategoricalDistribution([False, True]),
         }
diff --git a/etna/transforms/math/differencing.py b/etna/transforms/math/differencing.py
index 3f6f27317..f58f33cf0 100644
--- a/etna/transforms/math/differencing.py
+++ b/etna/transforms/math/differencing.py
@@ -8,16 +8,13 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.base import ReversibleTransform
 from etna.transforms.utils import check_new_segments
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class _SingleDifferencingTransform(ReversibleTransform):
     """Calculate a time series differences of order 1.
@@ -479,7 +476,7 @@ def _inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:
             result_df = transform._inverse_transform(result_df)
         return result_df
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``order`` parameter. Other parameters are expected to be set by the user.
@@ -490,5 +487,5 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "order": IntUniformDistribution(low=1, high=2),
+            "order": IntDistribution(low=1, high=2),
         }
diff --git a/etna/transforms/math/power.py b/etna/transforms/math/power.py
index 51ac53fe6..9fad2c461 100644
--- a/etna/transforms/math/power.py
+++ b/etna/transforms/math/power.py
@@ -5,14 +5,11 @@
 
 from sklearn.preprocessing import PowerTransformer
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.math.sklearn import SklearnTransform
 from etna.transforms.math.sklearn import TransformMode
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class YeoJohnsonTransform(SklearnTransform):
     """YeoJohnsonTransform applies Yeo-Johns transformation to a DataFrame.
@@ -64,7 +61,7 @@ def __init__(
             mode=mode,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``mode``, ``standardize``. Other parameters are expected to be set by the user.
@@ -133,7 +130,7 @@ def __init__(
             mode=mode,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``mode``, ``standardize``. Other parameters are expected to be set by the user.
diff --git a/etna/transforms/math/scalers.py b/etna/transforms/math/scalers.py
index 3ab6a55a3..776c19af1 100644
--- a/etna/transforms/math/scalers.py
+++ b/etna/transforms/math/scalers.py
@@ -9,14 +9,11 @@
 from sklearn.preprocessing import RobustScaler
 from sklearn.preprocessing import StandardScaler
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.math.sklearn import SklearnTransform
 from etna.transforms.math.sklearn import TransformMode
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class StandardScalerTransform(SklearnTransform):
     """Standardize features by removing the mean and scaling to unit variance.
@@ -75,7 +72,7 @@ def __init__(
             mode=mode,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``mode``, ``with_mean``, ``with_std``.
@@ -171,7 +168,7 @@ def __init__(
             mode=mode,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``mode``, ``with_centering``, ``with_scaling``, ``unit_variance``.
@@ -250,7 +247,7 @@ def __init__(
             mode=mode,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``mode``, ``clip``. Other parameters are expected to be set by the user.
diff --git a/etna/transforms/math/sklearn.py b/etna/transforms/math/sklearn.py
index 28742c61d..a3614f1d1 100644
--- a/etna/transforms/math/sklearn.py
+++ b/etna/transforms/math/sklearn.py
@@ -10,18 +10,15 @@
 import pandas as pd
 from sklearn.base import TransformerMixin
 
-from etna import SETTINGS
 from etna.core import StringEnumWithRepr
 from etna.datasets import TSDataset
 from etna.datasets import set_columns_wide
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import ReversibleTransform
 from etna.transforms.utils import check_new_segments
 from etna.transforms.utils import match_target_quantiles
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class TransformMode(StringEnumWithRepr):
     """Enum for different metric aggregation modes."""
@@ -304,7 +301,7 @@ def get_regressors_info(self) -> List[str]:
             return []
         return self.out_column_regressors
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes ``mode`` parameter. Other parameters are expected to be set by the user.
diff --git a/etna/transforms/math/statistics.py b/etna/transforms/math/statistics.py
index a4b5023b5..c04bba4e8 100644
--- a/etna/transforms/math/statistics.py
+++ b/etna/transforms/math/statistics.py
@@ -8,15 +8,12 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.base import IrreversibleTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import UniformDistribution
-
 
 class WindowStatisticsTransform(IrreversibleTransform, ABC):
     """WindowStatisticsTransform handles computation of statistical features on windows."""
@@ -119,7 +116,7 @@ def get_regressors_info(self) -> List[str]:
             raise ValueError("Fit the transform to get the correct regressors info!")
         return [self.out_column_name] if self.in_column_regressor else []
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes only ``window`` parameter. Other parameters are expected to be set by the user.
@@ -130,7 +127,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "window": IntUniformDistribution(low=1, high=20),
+            "window": IntDistribution(low=1, high=20),
         }
 
 
@@ -214,7 +211,7 @@ def _aggregate(self, series: np.ndarray) -> np.ndarray:
             mean[:, segment] = bn.nanmean(series[:, segment] * self._alpha_range, axis=1)
         return mean
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``window``, ``alpha``. Other parameters are expected to be set by the user.
@@ -227,7 +224,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         grid = super().params_to_tune()
         grid.update(
             {
-                "alpha": UniformDistribution(low=0.2, high=1),
+                "alpha": FloatDistribution(low=0.2, high=1),
             }
         )
         return grid
@@ -348,7 +345,7 @@ def _aggregate(self, series: np.ndarray) -> np.ndarray:
         series = np.apply_along_axis(np.nanquantile, axis=2, arr=series, q=self.quantile)
         return series
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``window``, ``quantile``. Other parameters are expected to be set by the user.
@@ -361,7 +358,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         grid = super().params_to_tune()
         grid.update(
             {
-                "quantile": UniformDistribution(low=0, high=1),
+                "quantile": FloatDistribution(low=0, high=1),
             }
         )
         return grid
diff --git a/etna/transforms/missing_values/imputation.py b/etna/transforms/missing_values/imputation.py
index f7fef6345..e33d85cc9 100644
--- a/etna/transforms/missing_values/imputation.py
+++ b/etna/transforms/missing_values/imputation.py
@@ -6,15 +6,12 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.base import OneSegmentTransform
 from etna.transforms.base import ReversiblePerSegmentWrapper
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-
 
 class ImputerMode(str, Enum):
     """Enum for different imputation strategy."""
@@ -284,7 +281,7 @@ def get_regressors_info(self) -> List[str]:
         """Return the list with regressors created by the transform."""
         return []
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``strategy``, ``window``.
@@ -300,12 +297,12 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
         if self.seasonality > 1:
             return {
                 "strategy": CategoricalDistribution(["constant", "mean", "running_mean", "forward_fill", "seasonal"]),
-                "window": IntUniformDistribution(low=1, high=20),
+                "window": IntDistribution(low=1, high=20),
             }
         else:
             return {
                 "strategy": CategoricalDistribution(["constant", "mean", "running_mean", "forward_fill"]),
-                "window": IntUniformDistribution(low=1, high=20),
+                "window": IntDistribution(low=1, high=20),
             }
 
 
diff --git a/etna/transforms/outliers/point_outliers.py b/etna/transforms/outliers/point_outliers.py
index 8506466b8..86e47ad7d 100644
--- a/etna/transforms/outliers/point_outliers.py
+++ b/etna/transforms/outliers/point_outliers.py
@@ -13,18 +13,16 @@
 from etna.analysis import get_anomalies_median
 from etna.analysis import get_anomalies_prediction_interval
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.models import SARIMAXModel
 from etna.transforms.outliers.base import OutliersTransform
 
 if SETTINGS.prophet_required:
     from etna.models import ProphetModel
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-    from optuna.distributions import IntUniformDistribution
-    from optuna.distributions import UniformDistribution
-
 
 class MedianOutliersTransform(OutliersTransform):
     """Transform that uses :py:func:`~etna.analysis.outliers.median_outliers.get_anomalies_median` to find anomalies in data.
@@ -66,7 +64,7 @@ def detect_outliers(self, ts: TSDataset) -> Dict[str, List[pd.Timestamp]]:
         """
         return get_anomalies_median(ts=ts, in_column=self.in_column, window_size=self.window_size, alpha=self.alpha)
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``window_size``, ``alpha``. Other parameters are expected to be set by the user.
@@ -77,8 +75,8 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "window_size": IntUniformDistribution(low=3, high=30),
-            "alpha": UniformDistribution(low=0.5, high=5),
+            "window_size": IntDistribution(low=3, high=30),
+            "alpha": FloatDistribution(low=0.5, high=5),
         }
 
 
@@ -142,7 +140,7 @@ def detect_outliers(self, ts: TSDataset) -> Dict[str, List[pd.Timestamp]]:
             distance_func=self.distance_func,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``window_size``, ``distance_coef``, ``n_neighbors``.
@@ -154,9 +152,9 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "window_size": IntUniformDistribution(low=3, high=30),
-            "distance_coef": UniformDistribution(low=0.5, high=5),
-            "n_neighbors": IntUniformDistribution(low=1, high=10),
+            "window_size": IntDistribution(low=3, high=30),
+            "distance_coef": FloatDistribution(low=0.5, high=5),
+            "n_neighbors": IntDistribution(low=1, high=10),
         }
 
 
@@ -223,7 +221,7 @@ def detect_outliers(self, ts: TSDataset) -> Dict[str, List[pd.Timestamp]]:
             **self.model_kwargs,
         )
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``interval_width``, ``model``. Other parameters are expected to be set by the user.
@@ -234,7 +232,7 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             Grid to tune.
         """
         return {
-            "interval_width": UniformDistribution(low=0.8, high=1.0),
+            "interval_width": FloatDistribution(low=0.8, high=1.0),
             "model": CategoricalDistribution(["prophet", "sarimax"]),
         }
 
diff --git a/etna/transforms/timestamp/date_flags.py b/etna/transforms/timestamp/date_flags.py
index 6a7a23afc..8f248641e 100644
--- a/etna/transforms/timestamp/date_flags.py
+++ b/etna/transforms/timestamp/date_flags.py
@@ -8,14 +8,11 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import FutureMixin
 from etna.transforms.base import IrreversibleTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class DateFlagsTransform(IrreversibleTransform, FutureMixin):
     """DateFlagsTransform is a class that implements extraction of the main date-based features from datetime column.
@@ -351,7 +348,7 @@ def _get_weekends(timestamp_series: pd.Series) -> np.ndarray:
         weekend_days = (5, 6)
         return timestamp_series.apply(lambda x: x.weekday() in weekend_days).values
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``day_number_in_week``, ``day_number_in_month``, ``day_number_in_year``,
diff --git a/etna/transforms/timestamp/fourier.py b/etna/transforms/timestamp/fourier.py
index d9d6ab4fb..cbc3f67a1 100644
--- a/etna/transforms/timestamp/fourier.py
+++ b/etna/transforms/timestamp/fourier.py
@@ -7,14 +7,11 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import IntDistribution
 from etna.transforms.base import FutureMixin
 from etna.transforms.base import IrreversibleTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import IntLogUniformDistribution
-
 
 class FourierTransform(IrreversibleTransform, FutureMixin):
     """Adds fourier features to the dataset.
@@ -160,7 +157,7 @@ def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
 
         return self._construct_answer(df, features)
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         If ``self.order`` is set then this grid tunes ``order`` parameter:
@@ -175,4 +172,4 @@ def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
             return {}
 
         max_value = math.ceil(self.period / 2)
-        return {"order": IntLogUniformDistribution(low=1, high=max_value)}
+        return {"order": IntDistribution(low=1, high=max_value, log=True)}
diff --git a/etna/transforms/timestamp/special_days.py b/etna/transforms/timestamp/special_days.py
index cc657151d..c4f38cc60 100644
--- a/etna/transforms/timestamp/special_days.py
+++ b/etna/transforms/timestamp/special_days.py
@@ -7,15 +7,12 @@
 
 import pandas as pd
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import FutureMixin
 from etna.transforms.base import IrreversiblePerSegmentWrapper
 from etna.transforms.base import OneSegmentTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 def calc_day_number_in_week(datetime_day: datetime.datetime) -> int:
     return datetime_day.weekday()
@@ -221,7 +218,7 @@ def get_regressors_info(self) -> List[str]:
             output_columns.append("anomaly_monthdays")
         return output_columns
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``find_special_weekday``, ``find_special_month_day``.
diff --git a/etna/transforms/timestamp/time_flags.py b/etna/transforms/timestamp/time_flags.py
index ac200d1e0..113d3ec0f 100644
--- a/etna/transforms/timestamp/time_flags.py
+++ b/etna/transforms/timestamp/time_flags.py
@@ -6,14 +6,11 @@
 import numpy as np
 import pandas as pd
 
-from etna import SETTINGS
+from etna.distributions import BaseDistribution
+from etna.distributions import CategoricalDistribution
 from etna.transforms.base import FutureMixin
 from etna.transforms.base import IrreversibleTransform
 
-if SETTINGS.auto_required:
-    from optuna.distributions import BaseDistribution
-    from optuna.distributions import CategoricalDistribution
-
 
 class TimeFlagsTransform(IrreversibleTransform, FutureMixin):
     """TimeFlagsTransform is a class that implements extraction of the main time-based features from datetime column."""
@@ -208,7 +205,7 @@ def _get_period_in_day(timestamp_series: pd.Series, period_in_hours: int = 12) -
         """
         return timestamp_series.apply(lambda x: x.hour // period_in_hours).values
 
-    def params_to_tune(self) -> Dict[str, "BaseDistribution"]:
+    def params_to_tune(self) -> Dict[str, BaseDistribution]:
         """Get default grid for tuning hyperparameters.
 
         This grid tunes parameters: ``minute_in_hour_number``, ``fifteen_minutes_in_hour_number``, ``hour_number``,
diff --git a/examples/automl.ipynb b/examples/automl.ipynb
index c1713521d..632d4c8e0 100644
--- a/examples/automl.ipynb
+++ b/examples/automl.ipynb
@@ -373,7 +373,9 @@
    "id": "8c732645",
    "metadata": {},
    "source": [
-    "Let's get back to our initial question about `params_to_tune`. In our optuna study we are going to sample each parameter value from its distribution and pass it into `pipeline.set_params` method. So, the keys for `params_to_tune` should be a valid for `set_params` method."
+    "Let's get back to our initial question about `params_to_tune`. In our optuna study we are going to sample each parameter value from its distribution and pass it into `pipeline.set_params` method. So, the keys for `params_to_tune` should be a valid for `set_params` method.\n",
+    "\n",
+    "Distributions are taken from `etna.distributions` and they are matching `optuna.Trial.suggest_` methods."
    ]
   },
   {
@@ -391,12 +393,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import optuna\n",
+    "from etna.distributions import CategoricalDistribution\n",
     "\n",
     "\n",
     "example_params_to_tune = {\n",
-    "    \"model.fit_intercept\": optuna.distributions.CategoricalDistribution([False, True]),\n",
-    "    \"transforms.0.is_weekend\": optuna.distributions.CategoricalDistribution([False, True]),\n",
+    "    \"model.fit_intercept\": CategoricalDistribution([False, True]),\n",
+    "    \"transforms.0.is_weekend\": CategoricalDistribution([False, True]),\n",
     "}"
    ]
   },
@@ -417,16 +419,16 @@
     {
      "data": {
       "text/plain": [
-       "{'model.fit_intercept': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.day_number_in_week': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.day_number_in_month': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.day_number_in_year': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.week_number_in_month': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.week_number_in_year': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.month_number_in_year': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.season_number': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.year_number': CategoricalDistribution(choices=(False, True)),\n",
-       " 'transforms.1.is_weekend': CategoricalDistribution(choices=(False, True))}"
+       "{'model.fit_intercept': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.day_number_in_week': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.day_number_in_month': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.day_number_in_year': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.week_number_in_month': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.week_number_in_year': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.month_number_in_year': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.season_number': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.year_number': CategoricalDistribution(choices=[False, True]),\n",
+       " 'transforms.1.is_weekend': CategoricalDistribution(choices=[False, True])}"
       ]
      },
      "execution_count": 10,
@@ -668,6 +670,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import optuna\n",
+    "\n",
+    "\n",
     "optuna.logging.set_verbosity(optuna.logging.CRITICAL)"
    ]
   },
@@ -2049,7 +2054,7 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>19</th>\n",
+       "      <th>1</th>\n",
        "      <td>af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.057438</td>\n",
@@ -2057,7 +2062,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>9</th>\n",
+       "      <th>14</th>\n",
        "      <td>d8215d95e2c6c9a4b4fdacf3fa77dddc</td>\n",
        "      <td>Pipeline(model = NaiveModel(lag = 7, ), transf...</td>\n",
        "      <td>5.164436</td>\n",
@@ -2065,7 +2070,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>12</th>\n",
+       "      <th>19</th>\n",
        "      <td>8f640faabcac0552153ca19337179f3b</td>\n",
        "      <td>Pipeline(model = HoltWintersModel(trend = 'add...</td>\n",
        "      <td>5.931951</td>\n",
@@ -2073,7 +2078,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>3</th>\n",
+       "      <th>2</th>\n",
        "      <td>d6a44adb551f1aec09ef37c14aed260f</td>\n",
        "      <td>Pipeline(model = SeasonalMovingAverageModel(wi...</td>\n",
        "      <td>6.197182</td>\n",
@@ -2081,7 +2086,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>18</th>\n",
+       "      <th>16</th>\n",
        "      <td>16eb77200eb2fd5dc1f6f2a5067884cd</td>\n",
        "      <td>Pipeline(model = HoltWintersModel(trend = 'add...</td>\n",
        "      <td>6.347734</td>\n",
@@ -2089,7 +2094,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>17</th>\n",
+       "      <th>3</th>\n",
        "      <td>4c07749e913403906cd033e4882fc4f9</td>\n",
        "      <td>Pipeline(model = SeasonalMovingAverageModel(wi...</td>\n",
        "      <td>6.529721</td>\n",
@@ -2097,7 +2102,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>5</th>\n",
+       "      <th>12</th>\n",
        "      <td>6e2eb71d033b6d0607f5b6d0a7596ce9</td>\n",
        "      <td>Pipeline(model = ProphetModel(growth = 'linear...</td>\n",
        "      <td>7.799984</td>\n",
@@ -2105,7 +2110,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>13</th>\n",
+       "      <th>7</th>\n",
        "      <td>6bb58e7ce09eab00448d5732240ec2ec</td>\n",
        "      <td>Pipeline(model = CatBoostMultiSegmentModel(ite...</td>\n",
        "      <td>7.814187</td>\n",
@@ -2113,7 +2118,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>1</th>\n",
+       "      <th>13</th>\n",
        "      <td>a640ddfb767ea0cbf31751ddda6e36ee</td>\n",
        "      <td>Pipeline(model = CatBoostMultiSegmentModel(ite...</td>\n",
        "      <td>7.816528</td>\n",
@@ -2121,7 +2126,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>10</th>\n",
+       "      <th>8</th>\n",
        "      <td>cfeb21bcf2e922a390ade8be9d845e0d</td>\n",
        "      <td>Pipeline(model = ProphetModel(growth = 'linear...</td>\n",
        "      <td>7.893421</td>\n",
@@ -2129,7 +2134,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>7</th>\n",
+       "      <th>18</th>\n",
        "      <td>a5e036978ef9cc9f297c9eb2c280af05</td>\n",
        "      <td>Pipeline(model = AutoARIMAModel(), transforms ...</td>\n",
        "      <td>8.297048</td>\n",
@@ -2137,7 +2142,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>16</th>\n",
+       "      <th>6</th>\n",
        "      <td>2e36e0b9cb67a43bb1bf96fa2ccf718f</td>\n",
        "      <td>Pipeline(model = LinearMultiSegmentModel(fit_i...</td>\n",
        "      <td>9.205423</td>\n",
@@ -2145,7 +2150,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>2</th>\n",
+       "      <th>11</th>\n",
        "      <td>8b9f5fa09754a80f17380dec2b998f1d</td>\n",
        "      <td>Pipeline(model = LinearPerSegmentModel(fit_int...</td>\n",
        "      <td>10.997462</td>\n",
@@ -2153,7 +2158,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>11</th>\n",
+       "      <th>17</th>\n",
        "      <td>d62c0579459d4a1b88aea8ed6effdf4e</td>\n",
        "      <td>Pipeline(model = MovingAverageModel(window = 1...</td>\n",
        "      <td>11.317256</td>\n",
@@ -2161,7 +2166,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>14</th>\n",
+       "      <th>0</th>\n",
        "      <td>5916e5b653295271c79caae490618ee9</td>\n",
        "      <td>Pipeline(model = MovingAverageModel(window = 2...</td>\n",
        "      <td>12.028916</td>\n",
@@ -2169,7 +2174,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>4</th>\n",
+       "      <th>15</th>\n",
        "      <td>5a91b6c8acc2c461913df44fd1429375</td>\n",
        "      <td>Pipeline(model = ElasticPerSegmentModel(alpha ...</td>\n",
        "      <td>12.213320</td>\n",
@@ -2177,7 +2182,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>6</th>\n",
+       "      <th>10</th>\n",
        "      <td>403b3e18012af5ff9815b408f5c2e47d</td>\n",
        "      <td>Pipeline(model = MovingAverageModel(window = 4...</td>\n",
        "      <td>12.243011</td>\n",
@@ -2185,7 +2190,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>15</th>\n",
+       "      <th>5</th>\n",
        "      <td>6cf8605e6c513053ac4f5203e330c59d</td>\n",
        "      <td>Pipeline(model = HoltWintersModel(trend = None...</td>\n",
        "      <td>15.473118</td>\n",
@@ -2193,7 +2198,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>8</th>\n",
+       "      <th>4</th>\n",
        "      <td>53e90ae4cf7f1f71e6396107549c25ef</td>\n",
        "      <td>Pipeline(model = NaiveModel(lag = 1, ), transf...</td>\n",
        "      <td>19.361078</td>\n",
@@ -2201,10 +2206,10 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>0</th>\n",
+       "      <th>9</th>\n",
        "      <td>90b31b54cb8c01867be05a3320852682</td>\n",
        "      <td>Pipeline(model = ElasticMultiSegmentModel(alph...</td>\n",
-       "      <td>35.971166</td>\n",
+       "      <td>35.971399</td>\n",
        "      <td>TrialState.COMPLETE</td>\n",
        "      <td>pool</td>\n",
        "    </tr>\n",
@@ -2214,70 +2219,70 @@
       ],
       "text/plain": [
        "                                hash  \\\n",
-       "19  af8088ac0abfde46e93a8dbb407a2117   \n",
-       "9   d8215d95e2c6c9a4b4fdacf3fa77dddc   \n",
-       "12  8f640faabcac0552153ca19337179f3b   \n",
-       "3   d6a44adb551f1aec09ef37c14aed260f   \n",
-       "18  16eb77200eb2fd5dc1f6f2a5067884cd   \n",
-       "17  4c07749e913403906cd033e4882fc4f9   \n",
-       "5   6e2eb71d033b6d0607f5b6d0a7596ce9   \n",
-       "13  6bb58e7ce09eab00448d5732240ec2ec   \n",
-       "1   a640ddfb767ea0cbf31751ddda6e36ee   \n",
-       "10  cfeb21bcf2e922a390ade8be9d845e0d   \n",
-       "7   a5e036978ef9cc9f297c9eb2c280af05   \n",
-       "16  2e36e0b9cb67a43bb1bf96fa2ccf718f   \n",
-       "2   8b9f5fa09754a80f17380dec2b998f1d   \n",
-       "11  d62c0579459d4a1b88aea8ed6effdf4e   \n",
-       "14  5916e5b653295271c79caae490618ee9   \n",
-       "4   5a91b6c8acc2c461913df44fd1429375   \n",
-       "6   403b3e18012af5ff9815b408f5c2e47d   \n",
-       "15  6cf8605e6c513053ac4f5203e330c59d   \n",
-       "8   53e90ae4cf7f1f71e6396107549c25ef   \n",
-       "0   90b31b54cb8c01867be05a3320852682   \n",
+       "1   af8088ac0abfde46e93a8dbb407a2117   \n",
+       "14  d8215d95e2c6c9a4b4fdacf3fa77dddc   \n",
+       "19  8f640faabcac0552153ca19337179f3b   \n",
+       "2   d6a44adb551f1aec09ef37c14aed260f   \n",
+       "16  16eb77200eb2fd5dc1f6f2a5067884cd   \n",
+       "3   4c07749e913403906cd033e4882fc4f9   \n",
+       "12  6e2eb71d033b6d0607f5b6d0a7596ce9   \n",
+       "7   6bb58e7ce09eab00448d5732240ec2ec   \n",
+       "13  a640ddfb767ea0cbf31751ddda6e36ee   \n",
+       "8   cfeb21bcf2e922a390ade8be9d845e0d   \n",
+       "18  a5e036978ef9cc9f297c9eb2c280af05   \n",
+       "6   2e36e0b9cb67a43bb1bf96fa2ccf718f   \n",
+       "11  8b9f5fa09754a80f17380dec2b998f1d   \n",
+       "17  d62c0579459d4a1b88aea8ed6effdf4e   \n",
+       "0   5916e5b653295271c79caae490618ee9   \n",
+       "15  5a91b6c8acc2c461913df44fd1429375   \n",
+       "10  403b3e18012af5ff9815b408f5c2e47d   \n",
+       "5   6cf8605e6c513053ac4f5203e330c59d   \n",
+       "4   53e90ae4cf7f1f71e6396107549c25ef   \n",
+       "9   90b31b54cb8c01867be05a3320852682   \n",
        "\n",
        "                                             pipeline  SMAPE_mean  \\\n",
-       "19  Pipeline(model = CatBoostPerSegmentModel(itera...    5.057438   \n",
-       "9   Pipeline(model = NaiveModel(lag = 7, ), transf...    5.164436   \n",
-       "12  Pipeline(model = HoltWintersModel(trend = 'add...    5.931951   \n",
-       "3   Pipeline(model = SeasonalMovingAverageModel(wi...    6.197182   \n",
-       "18  Pipeline(model = HoltWintersModel(trend = 'add...    6.347734   \n",
-       "17  Pipeline(model = SeasonalMovingAverageModel(wi...    6.529721   \n",
-       "5   Pipeline(model = ProphetModel(growth = 'linear...    7.799984   \n",
-       "13  Pipeline(model = CatBoostMultiSegmentModel(ite...    7.814187   \n",
-       "1   Pipeline(model = CatBoostMultiSegmentModel(ite...    7.816528   \n",
-       "10  Pipeline(model = ProphetModel(growth = 'linear...    7.893421   \n",
-       "7   Pipeline(model = AutoARIMAModel(), transforms ...    8.297048   \n",
-       "16  Pipeline(model = LinearMultiSegmentModel(fit_i...    9.205423   \n",
-       "2   Pipeline(model = LinearPerSegmentModel(fit_int...   10.997462   \n",
-       "11  Pipeline(model = MovingAverageModel(window = 1...   11.317256   \n",
-       "14  Pipeline(model = MovingAverageModel(window = 2...   12.028916   \n",
-       "4   Pipeline(model = ElasticPerSegmentModel(alpha ...   12.213320   \n",
-       "6   Pipeline(model = MovingAverageModel(window = 4...   12.243011   \n",
-       "15  Pipeline(model = HoltWintersModel(trend = None...   15.473118   \n",
-       "8   Pipeline(model = NaiveModel(lag = 1, ), transf...   19.361078   \n",
-       "0   Pipeline(model = ElasticMultiSegmentModel(alph...   35.971166   \n",
+       "1   Pipeline(model = CatBoostPerSegmentModel(itera...    5.057438   \n",
+       "14  Pipeline(model = NaiveModel(lag = 7, ), transf...    5.164436   \n",
+       "19  Pipeline(model = HoltWintersModel(trend = 'add...    5.931951   \n",
+       "2   Pipeline(model = SeasonalMovingAverageModel(wi...    6.197182   \n",
+       "16  Pipeline(model = HoltWintersModel(trend = 'add...    6.347734   \n",
+       "3   Pipeline(model = SeasonalMovingAverageModel(wi...    6.529721   \n",
+       "12  Pipeline(model = ProphetModel(growth = 'linear...    7.799984   \n",
+       "7   Pipeline(model = CatBoostMultiSegmentModel(ite...    7.814187   \n",
+       "13  Pipeline(model = CatBoostMultiSegmentModel(ite...    7.816528   \n",
+       "8   Pipeline(model = ProphetModel(growth = 'linear...    7.893421   \n",
+       "18  Pipeline(model = AutoARIMAModel(), transforms ...    8.297048   \n",
+       "6   Pipeline(model = LinearMultiSegmentModel(fit_i...    9.205423   \n",
+       "11  Pipeline(model = LinearPerSegmentModel(fit_int...   10.997462   \n",
+       "17  Pipeline(model = MovingAverageModel(window = 1...   11.317256   \n",
+       "0   Pipeline(model = MovingAverageModel(window = 2...   12.028916   \n",
+       "15  Pipeline(model = ElasticPerSegmentModel(alpha ...   12.213320   \n",
+       "10  Pipeline(model = MovingAverageModel(window = 4...   12.243011   \n",
+       "5   Pipeline(model = HoltWintersModel(trend = None...   15.473118   \n",
+       "4   Pipeline(model = NaiveModel(lag = 1, ), transf...   19.361078   \n",
+       "9   Pipeline(model = ElasticMultiSegmentModel(alph...   35.971399   \n",
        "\n",
        "                  state study  \n",
+       "1   TrialState.COMPLETE  pool  \n",
+       "14  TrialState.COMPLETE  pool  \n",
        "19  TrialState.COMPLETE  pool  \n",
-       "9   TrialState.COMPLETE  pool  \n",
-       "12  TrialState.COMPLETE  pool  \n",
+       "2   TrialState.COMPLETE  pool  \n",
+       "16  TrialState.COMPLETE  pool  \n",
        "3   TrialState.COMPLETE  pool  \n",
+       "12  TrialState.COMPLETE  pool  \n",
+       "7   TrialState.COMPLETE  pool  \n",
+       "13  TrialState.COMPLETE  pool  \n",
+       "8   TrialState.COMPLETE  pool  \n",
        "18  TrialState.COMPLETE  pool  \n",
+       "6   TrialState.COMPLETE  pool  \n",
+       "11  TrialState.COMPLETE  pool  \n",
        "17  TrialState.COMPLETE  pool  \n",
-       "5   TrialState.COMPLETE  pool  \n",
-       "13  TrialState.COMPLETE  pool  \n",
-       "1   TrialState.COMPLETE  pool  \n",
+       "0   TrialState.COMPLETE  pool  \n",
+       "15  TrialState.COMPLETE  pool  \n",
        "10  TrialState.COMPLETE  pool  \n",
-       "7   TrialState.COMPLETE  pool  \n",
-       "16  TrialState.COMPLETE  pool  \n",
-       "2   TrialState.COMPLETE  pool  \n",
-       "11  TrialState.COMPLETE  pool  \n",
-       "14  TrialState.COMPLETE  pool  \n",
+       "5   TrialState.COMPLETE  pool  \n",
        "4   TrialState.COMPLETE  pool  \n",
-       "6   TrialState.COMPLETE  pool  \n",
-       "15  TrialState.COMPLETE  pool  \n",
-       "8   TrialState.COMPLETE  pool  \n",
-       "0   TrialState.COMPLETE  pool  "
+       "9   TrialState.COMPLETE  pool  "
       ]
      },
      "execution_count": 25,
@@ -2354,7 +2359,7 @@
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
-       "      <th>56</th>\n",
+       "      <th>82</th>\n",
        "      <td>419fc80cf634ba0888c4f899f666ad45</td>\n",
        "      <td>Pipeline(model = HoltWintersModel(trend = 'mul...</td>\n",
        "      <td>4.769471</td>\n",
@@ -2362,7 +2367,7 @@
        "      <td>tuning/8f640faabcac0552153ca19337179f3b</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>89</th>\n",
+       "      <th>37</th>\n",
        "      <td>731ccb72a473bec81789b7f186001ddd</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>4.899715</td>\n",
@@ -2370,7 +2375,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>97</th>\n",
+       "      <th>45</th>\n",
        "      <td>9c302769456b4adb9143f11c582f7264</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>4.927197</td>\n",
@@ -2378,7 +2383,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>88</th>\n",
+       "      <th>36</th>\n",
        "      <td>182c748af70287ab3a12bf32c03320f5</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>4.941247</td>\n",
@@ -2386,7 +2391,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>96</th>\n",
+       "      <th>44</th>\n",
        "      <td>4f426335c0eb00d847d9dd1e0a421415</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>4.977773</td>\n",
@@ -2394,7 +2399,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>98</th>\n",
+       "      <th>46</th>\n",
        "      <td>2cafd0750f191e7ab2d4160da50a7c64</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.056993</td>\n",
@@ -2402,7 +2407,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>19</th>\n",
+       "      <th>1</th>\n",
        "      <td>af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.057438</td>\n",
@@ -2410,7 +2415,7 @@
        "      <td>pool</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>75</th>\n",
+       "      <th>23</th>\n",
        "      <td>382825866425cac211691205a9537c95</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.081609</td>\n",
@@ -2418,7 +2423,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>95</th>\n",
+       "      <th>43</th>\n",
        "      <td>c2a8d498fe35873d060e173e1af042d5</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.117583</td>\n",
@@ -2426,7 +2431,7 @@
        "      <td>tuning/af8088ac0abfde46e93a8dbb407a2117</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>91</th>\n",
+       "      <th>39</th>\n",
        "      <td>035f8e28180bc7491a30b3d0d67060c9</td>\n",
        "      <td>Pipeline(model = CatBoostPerSegmentModel(itera...</td>\n",
        "      <td>5.135956</td>\n",
@@ -2439,40 +2444,40 @@
       ],
       "text/plain": [
        "                                hash  \\\n",
-       "56  419fc80cf634ba0888c4f899f666ad45   \n",
-       "89  731ccb72a473bec81789b7f186001ddd   \n",
-       "97  9c302769456b4adb9143f11c582f7264   \n",
-       "88  182c748af70287ab3a12bf32c03320f5   \n",
-       "96  4f426335c0eb00d847d9dd1e0a421415   \n",
-       "98  2cafd0750f191e7ab2d4160da50a7c64   \n",
-       "19  af8088ac0abfde46e93a8dbb407a2117   \n",
-       "75  382825866425cac211691205a9537c95   \n",
-       "95  c2a8d498fe35873d060e173e1af042d5   \n",
-       "91  035f8e28180bc7491a30b3d0d67060c9   \n",
+       "82  419fc80cf634ba0888c4f899f666ad45   \n",
+       "37  731ccb72a473bec81789b7f186001ddd   \n",
+       "45  9c302769456b4adb9143f11c582f7264   \n",
+       "36  182c748af70287ab3a12bf32c03320f5   \n",
+       "44  4f426335c0eb00d847d9dd1e0a421415   \n",
+       "46  2cafd0750f191e7ab2d4160da50a7c64   \n",
+       "1   af8088ac0abfde46e93a8dbb407a2117   \n",
+       "23  382825866425cac211691205a9537c95   \n",
+       "43  c2a8d498fe35873d060e173e1af042d5   \n",
+       "39  035f8e28180bc7491a30b3d0d67060c9   \n",
        "\n",
        "                                             pipeline  SMAPE_mean  \\\n",
-       "56  Pipeline(model = HoltWintersModel(trend = 'mul...    4.769471   \n",
-       "89  Pipeline(model = CatBoostPerSegmentModel(itera...    4.899715   \n",
-       "97  Pipeline(model = CatBoostPerSegmentModel(itera...    4.927197   \n",
-       "88  Pipeline(model = CatBoostPerSegmentModel(itera...    4.941247   \n",
-       "96  Pipeline(model = CatBoostPerSegmentModel(itera...    4.977773   \n",
-       "98  Pipeline(model = CatBoostPerSegmentModel(itera...    5.056993   \n",
-       "19  Pipeline(model = CatBoostPerSegmentModel(itera...    5.057438   \n",
-       "75  Pipeline(model = CatBoostPerSegmentModel(itera...    5.081609   \n",
-       "95  Pipeline(model = CatBoostPerSegmentModel(itera...    5.117583   \n",
-       "91  Pipeline(model = CatBoostPerSegmentModel(itera...    5.135956   \n",
+       "82  Pipeline(model = HoltWintersModel(trend = 'mul...    4.769471   \n",
+       "37  Pipeline(model = CatBoostPerSegmentModel(itera...    4.899715   \n",
+       "45  Pipeline(model = CatBoostPerSegmentModel(itera...    4.927197   \n",
+       "36  Pipeline(model = CatBoostPerSegmentModel(itera...    4.941247   \n",
+       "44  Pipeline(model = CatBoostPerSegmentModel(itera...    4.977773   \n",
+       "46  Pipeline(model = CatBoostPerSegmentModel(itera...    5.056993   \n",
+       "1   Pipeline(model = CatBoostPerSegmentModel(itera...    5.057438   \n",
+       "23  Pipeline(model = CatBoostPerSegmentModel(itera...    5.081609   \n",
+       "43  Pipeline(model = CatBoostPerSegmentModel(itera...    5.117583   \n",
+       "39  Pipeline(model = CatBoostPerSegmentModel(itera...    5.135956   \n",
        "\n",
        "                  state                                    study  \n",
-       "56  TrialState.COMPLETE  tuning/8f640faabcac0552153ca19337179f3b  \n",
-       "89  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "97  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "88  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "96  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "98  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "19  TrialState.COMPLETE                                     pool  \n",
-       "75  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "95  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
-       "91  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  "
+       "82  TrialState.COMPLETE  tuning/8f640faabcac0552153ca19337179f3b  \n",
+       "37  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "45  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "36  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "44  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "46  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "1   TrialState.COMPLETE                                     pool  \n",
+       "23  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "43  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  \n",
+       "39  TrialState.COMPLETE  tuning/af8088ac0abfde46e93a8dbb407a2117  "
       ]
      },
      "execution_count": 27,


Test patch:
diff --git a/tests/test_auto/__init__.py b/tests/test_auto/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/test_auto/optuna/test_config_sampler.py b/tests/test_auto/test_optuna/test_config_sampler.py
similarity index 100%
rename from tests/test_auto/optuna/test_config_sampler.py
rename to tests/test_auto/test_optuna/test_config_sampler.py
diff --git a/tests/test_auto/optuna/test_wrapper.py b/tests/test_auto/test_optuna/test_wrapper.py
similarity index 100%
rename from tests/test_auto/optuna/test_wrapper.py
rename to tests/test_auto/test_optuna/test_wrapper.py
diff --git a/tests/test_auto/pool/test_pool.py b/tests/test_auto/test_pool/test_pool.py
similarity index 100%
rename from tests/test_auto/pool/test_pool.py
rename to tests/test_auto/test_pool/test_pool.py
diff --git a/tests/test_auto/runner/test_local.py b/tests/test_auto/test_runner/test_local.py
similarity index 100%
rename from tests/test_auto/runner/test_local.py
rename to tests/test_auto/test_runner/test_local.py
diff --git a/tests/test_auto/test_tune.py b/tests/test_auto/test_tune.py
index f45adbc3b..0a1b972a0 100644
--- a/tests/test_auto/test_tune.py
+++ b/tests/test_auto/test_tune.py
@@ -4,17 +4,14 @@
 
 import pandas as pd
 import pytest
-from optuna.distributions import CategoricalDistribution
-from optuna.distributions import DiscreteUniformDistribution
-from optuna.distributions import IntLogUniformDistribution
-from optuna.distributions import IntUniformDistribution
-from optuna.distributions import LogUniformDistribution
-from optuna.distributions import UniformDistribution
 from typing_extensions import Literal
 
 from etna.auto import Tune
 from etna.auto.auto import _Callback
 from etna.auto.auto import _Initializer
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.metrics import MAE
 from etna.models import NaiveModel
 from etna.models import SimpleExpSmoothingModel
@@ -98,13 +95,14 @@ def test_init_optuna(
 @pytest.mark.parametrize(
     "params, model",
     [
-        ({"model.smoothing_level": UniformDistribution(0.1, 1)}, SimpleExpSmoothingModel()),
-        ({"model.smoothing_level": LogUniformDistribution(0.1, 1)}, SimpleExpSmoothingModel()),
-        ({"model.smoothing_level": DiscreteUniformDistribution(0.1, 1, 0.1)}, SimpleExpSmoothingModel()),
-        ({"model.lag": IntUniformDistribution(1, 5)}, NaiveModel()),
-        ({"model.lag": IntLogUniformDistribution(1, 5)}, NaiveModel()),
+        ({"model.smoothing_level": FloatDistribution(low=0.1, high=1)}, SimpleExpSmoothingModel()),
+        ({"model.smoothing_level": FloatDistribution(low=0.1, high=1, log=True)}, SimpleExpSmoothingModel()),
+        ({"model.smoothing_level": FloatDistribution(low=0.1, high=1, step=0.1)}, SimpleExpSmoothingModel()),
+        ({"model.lag": IntDistribution(low=1, high=5)}, NaiveModel()),
+        ({"model.lag": IntDistribution(low=1, high=5, log=True)}, NaiveModel()),
+        ({"model.lag": IntDistribution(low=1, high=5, step=2)}, NaiveModel()),
         ({"model.lag": CategoricalDistribution((1, 2, 3))}, NaiveModel()),
-        ({"model.smoothing_level": UniformDistribution(1, 5)}, SimpleExpSmoothingModel()),
+        ({"model.smoothing_level": FloatDistribution(low=1, high=5)}, SimpleExpSmoothingModel()),
     ],
 )
 def test_can_handle_distribution_type(example_tsds, optuna_storage, params, model):
@@ -115,7 +113,10 @@ def test_can_handle_distribution_type(example_tsds, optuna_storage, params, mode
 
 
 def test_can_handle_transforms(example_tsds, optuna_storage):
-    params = {"transforms.0.value": IntUniformDistribution(0, 17), "transforms.1.value": IntUniformDistribution(0, 17)}
+    params = {
+        "transforms.0.value": IntDistribution(low=0, high=17),
+        "transforms.1.value": IntDistribution(low=0, high=17),
+    }
     with patch.object(Pipeline, "params_to_tune", return_value=params):
         pipeline = Pipeline(
             NaiveModel(),
@@ -191,7 +192,7 @@ def test_tune_run(example_tsds, optuna_storage, pipeline):
 @pytest.mark.parametrize(
     "pipeline, params_to_tune",
     [
-        (Pipeline(NaiveModel(1), horizon=7), {"model.lag": IntUniformDistribution(low=1, high=5)}),
+        (Pipeline(NaiveModel(1), horizon=7), {"model.lag": IntDistribution(low=1, high=5)}),
     ],
 )
 def test_tune_run_custom_params_to_tune(example_tsds, optuna_storage, pipeline, params_to_tune):
diff --git a/tests/test_auto/test_utils.py b/tests/test_auto/test_utils.py
new file mode 100644
index 000000000..b3a582076
--- /dev/null
+++ b/tests/test_auto/test_utils.py
@@ -0,0 +1,46 @@
+from dataclasses import asdict
+from unittest.mock import Mock
+
+import pytest
+
+from etna.auto.utils import suggest_parameters
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
+
+
+@pytest.mark.parametrize(
+    "params_to_tune",
+    [
+        {"x": CategoricalDistribution([False, True])},
+        {"x": CategoricalDistribution(["A", "B", "C"])},
+        {"x": IntDistribution(low=1, high=5)},
+        {"x": IntDistribution(low=1, high=5, step=2)},
+        {"x": IntDistribution(low=1, high=5, log=True)},
+        {"x": FloatDistribution(low=0.1, high=10.0)},
+        {"x": FloatDistribution(low=0.1, high=10.0, step=0.1)},
+        {"x": FloatDistribution(low=1e-2, high=100, log=True)},
+        {
+            "c1": CategoricalDistribution([False, True]),
+            "c2": CategoricalDistribution(["A", "B", "C"]),
+            "i1": IntDistribution(low=1, high=5),
+            "i2": IntDistribution(low=1, high=5, step=2),
+            "f1": FloatDistribution(low=0.1, high=10.0),
+            "f2": FloatDistribution(low=0.1, high=10.0, step=0.1),
+        },
+    ],
+)
+def test_suggest_parameters(params_to_tune):
+    trial = Mock()
+    method_names = {
+        CategoricalDistribution: "suggest_categorical",
+        IntDistribution: "suggest_int",
+        FloatDistribution: "suggest_float",
+    }
+
+    results = suggest_parameters(trial=trial, params_to_tune=params_to_tune)
+
+    assert results.keys() == params_to_tune.keys()
+    for param_name in results.keys():
+        method = getattr(trial, method_names[type(params_to_tune[param_name])])
+        method.assert_any_call(param_name, **asdict(params_to_tune[param_name]))
diff --git a/tests/test_models/test_holt_winters_model.py b/tests/test_models/test_holt_winters_model.py
index 6e88e3c55..8fb2c79f6 100644
--- a/tests/test_models/test_holt_winters_model.py
+++ b/tests/test_models/test_holt_winters_model.py
@@ -335,6 +335,16 @@ def test_prediction_decomposition(outliers_tsds, model):
     ],
 )
 def test_params_to_tune(model, expected_length, const_ts):
+    def skip_parameters(parameters):
+        if (
+            "damped_trend" in parameters
+            and "trend" in parameters
+            and parameters["damped_trend"]
+            and parameters["trend"] is None
+        ):
+            return True
+        return False
+
     ts = const_ts
     assert len(model.params_to_tune()) == expected_length
-    assert_sampling_is_valid(model=model, ts=ts)
+    assert_sampling_is_valid(model=model, ts=ts, skip_parameters=skip_parameters)
diff --git a/tests/test_models/utils.py b/tests/test_models/utils.py
index fac451293..c1cbbd830 100644
--- a/tests/test_models/utils.py
+++ b/tests/test_models/utils.py
@@ -1,12 +1,15 @@
 import pathlib
 import tempfile
+from typing import Callable
+from typing import Optional
 from typing import Sequence
 from typing import Tuple
 
+import optuna
 import pandas as pd
 from lightning_fabric.utilities.seed import seed_everything
-from optuna.samplers import RandomSampler
 
+from etna.auto.utils import suggest_parameters
 from etna.datasets import TSDataset
 from etna.models.base import ModelType
 from etna.pipeline import Pipeline
@@ -41,14 +44,20 @@ def assert_model_equals_loaded_original(
     return model, loaded_model
 
 
-def assert_sampling_is_valid(model: ModelType, ts: TSDataset, seed: int = 0):
-    grid = model.params_to_tune()
-    # we need sampler to get a value from distribution
-    sampler = RandomSampler(seed=seed)
-    for name, distribution in grid.items():
-        value = sampler.sample_independent(study=None, trial=None, param_name=name, param_distribution=distribution)
-        new_model = model.set_params(**{name: value})
-        new_model.fit(ts)
+def assert_sampling_is_valid(
+    model: ModelType, ts: TSDataset, seed: int = 0, n_trials: int = 5, skip_parameters: Optional[Callable] = None
+):
+    params_to_tune = model.params_to_tune()
+
+    def _objective(trial: optuna.Trial) -> float:
+        parameters = suggest_parameters(trial, params_to_tune)
+        if skip_parameters is None or not skip_parameters(parameters):
+            new_model = model.set_params(**parameters)
+            new_model.fit(ts)
+        return 0.0
+
+    study = optuna.create_study(sampler=optuna.samplers.RandomSampler(seed=seed))
+    study.optimize(_objective, n_trials=n_trials)
 
 
 def assert_prediction_components_are_present(model, train, test, prediction_size=None):
diff --git a/tests/test_pipeline/test_autoregressive_pipeline.py b/tests/test_pipeline/test_autoregressive_pipeline.py
index bb71f38ce..8acf99e19 100644
--- a/tests/test_pipeline/test_autoregressive_pipeline.py
+++ b/tests/test_pipeline/test_autoregressive_pipeline.py
@@ -7,11 +7,11 @@
 import numpy as np
 import pandas as pd
 import pytest
-from optuna.distributions import CategoricalDistribution
-from optuna.distributions import IntUniformDistribution
-from optuna.distributions import LogUniformDistribution
 
 from etna.datasets import TSDataset
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.metrics import MAE
 from etna.metrics import MetricAggregationMode
 from etna.models import CatBoostMultiSegmentModel
@@ -402,10 +402,10 @@ def test_predict_return_components(
             CatBoostMultiSegmentModel(iterations=100),
             [DateFlagsTransform(), LagTransform(in_column="target", lags=list(range(3, 10)))],
             {
-                "model.learning_rate": LogUniformDistribution(low=1e-4, high=0.5),
-                "model.depth": IntUniformDistribution(low=1, high=11, step=1),
-                "model.l2_leaf_reg": LogUniformDistribution(low=0.1, high=200.0),
-                "model.random_strength": LogUniformDistribution(low=1e-05, high=10.0),
+                "model.learning_rate": FloatDistribution(low=1e-4, high=0.5, log=True),
+                "model.depth": IntDistribution(low=1, high=11, step=1),
+                "model.l2_leaf_reg": FloatDistribution(low=0.1, high=200.0, log=True),
+                "model.random_strength": FloatDistribution(low=1e-05, high=10.0, log=True),
                 "transforms.0.day_number_in_week": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_month": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_year": CategoricalDistribution([False, True]),
diff --git a/tests/test_pipeline/test_base.py b/tests/test_pipeline/test_base.py
index 99a85fff5..462df4022 100644
--- a/tests/test_pipeline/test_base.py
+++ b/tests/test_pipeline/test_base.py
@@ -5,9 +5,9 @@
 
 import pandas as pd
 import pytest
-from optuna.distributions import BaseDistribution
 
 from etna.datasets import TSDataset
+from etna.distributions import BaseDistribution
 from etna.pipeline.base import BasePipeline
 
 
diff --git a/tests/test_pipeline/test_hierarchical_pipeline.py b/tests/test_pipeline/test_hierarchical_pipeline.py
index f0651d36f..c71ee40cc 100644
--- a/tests/test_pipeline/test_hierarchical_pipeline.py
+++ b/tests/test_pipeline/test_hierarchical_pipeline.py
@@ -5,12 +5,12 @@
 import numpy as np
 import pandas as pd
 import pytest
-from optuna.distributions import CategoricalDistribution
-from optuna.distributions import IntUniformDistribution
-from optuna.distributions import LogUniformDistribution
 
 from etna.datasets import TSDataset
 from etna.datasets.utils import match_target_quantiles
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.metrics import MAE
 from etna.metrics import Coverage
 from etna.metrics import Width
@@ -645,10 +645,10 @@ def test_predict_with_return_components(
             CatBoostMultiSegmentModel(iterations=100),
             [DateFlagsTransform(), LagTransform(in_column="target", lags=list(range(3, 10)))],
             {
-                "model.learning_rate": LogUniformDistribution(low=1e-4, high=0.5),
-                "model.depth": IntUniformDistribution(low=1, high=11, step=1),
-                "model.l2_leaf_reg": LogUniformDistribution(low=0.1, high=200.0),
-                "model.random_strength": LogUniformDistribution(low=1e-05, high=10.0),
+                "model.learning_rate": FloatDistribution(low=1e-4, high=0.5, log=True),
+                "model.depth": IntDistribution(low=1, high=11, step=1),
+                "model.l2_leaf_reg": FloatDistribution(low=0.1, high=200.0, log=True),
+                "model.random_strength": FloatDistribution(low=1e-05, high=10.0, log=True),
                 "transforms.0.day_number_in_week": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_month": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_year": CategoricalDistribution([False, True]),
diff --git a/tests/test_pipeline/test_pipeline.py b/tests/test_pipeline/test_pipeline.py
index b6021d470..8f75b406e 100644
--- a/tests/test_pipeline/test_pipeline.py
+++ b/tests/test_pipeline/test_pipeline.py
@@ -7,12 +7,12 @@
 import numpy as np
 import pandas as pd
 import pytest
-from optuna.distributions import CategoricalDistribution
-from optuna.distributions import IntUniformDistribution
-from optuna.distributions import LogUniformDistribution
 
 from etna.datasets import TSDataset
 from etna.datasets import generate_ar_df
+from etna.distributions import CategoricalDistribution
+from etna.distributions import FloatDistribution
+from etna.distributions import IntDistribution
 from etna.metrics import MAE
 from etna.metrics import MSE
 from etna.metrics import SMAPE
@@ -1218,10 +1218,10 @@ def test_predict_return_components(
             CatBoostMultiSegmentModel(iterations=100),
             [DateFlagsTransform(), LagTransform(in_column="target", lags=list(range(3, 10)))],
             {
-                "model.learning_rate": LogUniformDistribution(low=1e-4, high=0.5),
-                "model.depth": IntUniformDistribution(low=1, high=11, step=1),
-                "model.l2_leaf_reg": LogUniformDistribution(low=0.1, high=200.0),
-                "model.random_strength": LogUniformDistribution(low=1e-05, high=10.0),
+                "model.learning_rate": FloatDistribution(low=1e-4, high=0.5, log=True),
+                "model.depth": IntDistribution(low=1, high=11, step=1),
+                "model.l2_leaf_reg": FloatDistribution(low=0.1, high=200.0, log=True),
+                "model.random_strength": FloatDistribution(low=1e-05, high=10.0, log=True),
                 "transforms.0.day_number_in_week": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_month": CategoricalDistribution([False, True]),
                 "transforms.0.day_number_in_year": CategoricalDistribution([False, True]),
diff --git a/tests/test_transforms/test_timestamp/test_dateflags_transform.py b/tests/test_transforms/test_timestamp/test_dateflags_transform.py
index 7a351a8e0..1d08a25e9 100644
--- a/tests/test_transforms/test_timestamp/test_dateflags_transform.py
+++ b/tests/test_transforms/test_timestamp/test_dateflags_transform.py
@@ -293,7 +293,23 @@ def test_save_load(train_ts):
 
 
 def test_params_to_tune(train_ts):
+    def skip_parameters(parameters):
+        names = [
+            "day_number_in_week",
+            "day_number_in_month",
+            "day_number_in_year",
+            "week_number_in_year",
+            "week_number_in_month",
+            "month_number_in_year",
+            "season_number",
+            "year_number",
+        ]
+        values = [not parameters[x] for x in names]
+        if all(values):
+            return True
+        return False
+
     transform = DateFlagsTransform()
     ts = train_ts
     assert len(transform.params_to_tune()) > 0
-    assert_sampling_is_valid(transform=transform, ts=ts)
+    assert_sampling_is_valid(transform=transform, ts=ts, skip_parameters=skip_parameters)
diff --git a/tests/test_transforms/test_timestamp/test_special_days_transform.py b/tests/test_transforms/test_timestamp/test_special_days_transform.py
index 1925f6e75..e10565aa4 100644
--- a/tests/test_transforms/test_timestamp/test_special_days_transform.py
+++ b/tests/test_transforms/test_timestamp/test_special_days_transform.py
@@ -211,7 +211,14 @@ def test_save_load(ts_with_specials):
 
 
 def test_params_to_tune(ts_with_specials):
+    def skip_parameters(parameters):
+        names = ["find_special_weekday", "find_special_month_day"]
+        values = [not parameters[x] for x in names]
+        if all(values):
+            return True
+        return False
+
     transform = SpecialDaysTransform()
     ts = ts_with_specials
     assert len(transform.params_to_tune()) > 0
-    assert_sampling_is_valid(transform=transform, ts=ts)
+    assert_sampling_is_valid(transform=transform, ts=ts, skip_parameters=skip_parameters)
diff --git a/tests/test_transforms/test_timestamp/test_timeflags_transform.py b/tests/test_transforms/test_timestamp/test_timeflags_transform.py
index 592d65c5e..02caace68 100644
--- a/tests/test_transforms/test_timestamp/test_timeflags_transform.py
+++ b/tests/test_transforms/test_timestamp/test_timeflags_transform.py
@@ -227,7 +227,21 @@ def test_save_load(train_ts):
 
 
 def test_params_to_tune(train_ts):
+    def skip_parameters(parameters):
+        names = [
+            "minute_in_hour_number",
+            "fifteen_minutes_in_hour_number",
+            "hour_number",
+            "half_hour_number",
+            "half_day_number",
+            "one_third_day_number",
+        ]
+        values = [not parameters[x] for x in names]
+        if all(values):
+            return True
+        return False
+
     transform = TimeFlagsTransform()
     ts = train_ts
     assert len(transform.params_to_tune()) > 0
-    assert_sampling_is_valid(transform=transform, ts=ts)
+    assert_sampling_is_valid(transform=transform, ts=ts, skip_parameters=skip_parameters)
diff --git a/tests/test_transforms/utils.py b/tests/test_transforms/utils.py
index cc74d0fc8..a7f63fbe2 100644
--- a/tests/test_transforms/utils.py
+++ b/tests/test_transforms/utils.py
@@ -1,11 +1,14 @@
 import pathlib
 import tempfile
 from copy import deepcopy
+from typing import Callable
+from typing import Optional
 from typing import Tuple
 
+import optuna
 import pandas as pd
-from optuna.samplers import RandomSampler
 
+from etna.auto.utils import suggest_parameters
 from etna.datasets import TSDataset
 from etna.transforms import Transform
 
@@ -33,11 +36,17 @@ def assert_transformation_equals_loaded_original(transform: Transform, ts: TSDat
     return transform, loaded_transform
 
 
-def assert_sampling_is_valid(transform: Transform, ts: TSDataset, seed: int = 0):
-    grid = transform.params_to_tune()
-    # we need sampler to get a value from distribution
-    sampler = RandomSampler(seed=seed)
-    for name, distribution in grid.items():
-        value = sampler.sample_independent(study=None, trial=None, param_name=name, param_distribution=distribution)
-        new_transform = transform.set_params(**{name: value})
-        new_transform.fit(ts)
+def assert_sampling_is_valid(
+    transform: Transform, ts: TSDataset, seed: int = 0, n_trials: int = 3, skip_parameters: Optional[Callable] = None
+):
+    params_to_tune = transform.params_to_tune()
+
+    def _objective(trial: optuna.Trial) -> float:
+        parameters = suggest_parameters(trial, params_to_tune)
+        if skip_parameters is None or not skip_parameters(parameters):
+            new_transform = transform.set_params(**parameters)
+            new_transform.fit(ts)
+        return 0.0
+
+    study = optuna.create_study(sampler=optuna.samplers.RandomSampler(seed=seed))
+    study.optimize(_objective, n_trials=n_trials)