Translate input_size_multiplier in user-supplied Auto configs (#1553)

Co-authored-by: Marco <marco@nixtla.io>

Author: Vedant-Agarwal

neuralforecast/common/_base_auto.py

@@ -198,6 +198,10 @@ def __init__(
             raise ValueError(
                 f"Unknown backend {backend}. The supported backends are 'ray' and 'optuna'."
             )
+        # Translate `*input_size_multiplier` entries (as used by the models'
+        # `default_config`) into concrete `*input_size` ones, so that configs
+        # built on top of `default_config` are valid `config` arguments.
+        config_base = self._translate_input_size_multipliers(config_base, h)
 
         # Shallow-copy user-supplied options so subsequent mutations
         # (default resolution) don't leak back to the caller.
@@ -268,7 +272,9 @@ def __init__(
         else:
 
             def config_f(trial):
-                return {**config(trial), **config_base}
+                return self._translate_input_size_multipliers(
+                    {**config(trial), **config_base}, h
+                )
 
             self.config = config_f
 
@@ -300,6 +306,43 @@ def config_f(trial):
     def __repr__(self):
         return type(self).__name__ if self.alias is None else self.alias
 
+    @staticmethod
+    def _translate_input_size_multipliers(config, h):
+        """Translate `*input_size_multiplier` config entries into `*input_size` ones.
+
+        The models' `default_config` express the input sizes as multiples of the
+        horizon (`input_size_multiplier` and `inference_input_size_multiplier`).
+        These keys are not valid model arguments, so apply here the same
+        translation that `get_default_config` performs, which allows users to
+        provide configs built on top of `default_config`.
+
+        Args:
+            config (dict): Configuration dict, possibly containing
+                `*input_size_multiplier` entries.
+            h (int): Forecast horizon.
+
+        Returns:
+            dict: Configuration dict with the multipliers translated into sizes.
+        """
+        translations = {
+            "input_size_multiplier": "input_size",
+            "inference_input_size_multiplier": "inference_input_size",
+        }
+        if not any(key in config for key in translations):
+            return config
+        config = dict(config)
+        for multiplier_key, size_key in translations.items():
+            if multiplier_key not in config:
+                continue
+            multipliers = config.pop(multiplier_key)
+            if size_key in config:
+                continue
+            if isinstance(multipliers, (list, tuple)):
+                config[size_key] = tune.choice([h * x for x in multipliers])
+            else:
+                config[size_key] = h * multipliers
+        return config
+
     def _train_tune(self, config_step, cls_model, dataset, val_size, test_size):
         """BaseAuto._train_tune
 

tests/test_common/test_base_auto.py

@@ -166,6 +166,59 @@ def config_fn(trial):
         AutoMLP(h=4, config=config_fn, backend="optuna")
 
 
+def test_default_config_copy_is_valid_user_config(setup_module):
+    # https://github.com/Nixtla/neuralforecast/issues/571
+    # The default configs express the input size as `input_size_multiplier`,
+    # which used to be translated into `input_size` only when `config=None`,
+    # making configs built on top of `default_config` invalid.
+    dataset, _, _ = setup_module
+    config = AutoMLP.default_config.copy()
+    # tweak the search space, keeping it small so the test runs fast
+    config["input_size_multiplier"] = [1, 2]
+    config["hidden_size"] = tune.choice([8])
+    config["num_layers"] = 2
+    config["max_steps"] = 1
+    config["val_check_steps"] = 1
+    auto = AutoMLP(
+        h=12,
+        config=config,
+        num_samples=1,
+        ray_options=RayOptions(cpus=1, gpus=0),
+    )
+    assert "input_size_multiplier" not in auto.config
+    assert "input_size" in auto.config
+    auto.fit(dataset=dataset)
+    y_hat = auto.predict(dataset=dataset)
+    assert y_hat.shape[0] > 0
+
+
+def test_optuna_config_translates_input_size_multiplier(setup_module):
+    # https://github.com/Nixtla/neuralforecast/issues/571 (optuna variant)
+    dataset, _, _ = setup_module
+
+    def config_multiplier(trial):
+        return {
+            "input_size_multiplier": trial.suggest_categorical(
+                "input_size_multiplier", [1, 2]
+            ),
+            "hidden_size": trial.suggest_categorical("hidden_size", [8]),
+            "num_layers": 2,
+            "max_steps": 1,
+            "val_check_steps": 1,
+        }
+
+    auto = AutoMLP(
+        h=12,
+        config=config_multiplier,
+        backend="optuna",
+        search_alg=optuna.samplers.RandomSampler(seed=0),
+        num_samples=1,
+    )
+    auto.fit(dataset=dataset)
+    y_hat = auto.predict(dataset=dataset)
+    assert y_hat.shape[0] > 0
+
+
 def test_ray_time_budget(setup_module):
     dataset, _, _ = setup_module
     config = {