Build the mixed estimators from a typed config

sbi/neural_nets/factory.py

@@ -9,13 +9,12 @@
 
 from sbi.neural_nets.net_builders.estimator_configs import (
     VF_MODELS,
-    _BUILD_KWARG_ALIASES,
     _CLASSIFIER_CONFIGS,
     _DENSITY_CONFIGS,
-    ConditionalFlowConfig,
     MarginalFlowConfig,
     _config_from_factory_kwargs,
     _factory_defaults,
+    _mixed_config_from_factory_kwargs,
 )
 from sbi.neural_nets.net_builders.flow import (
     build_made,
@@ -139,34 +138,11 @@ def _density_family_args(embedding_net: nn.Module, **family_args: Any) -> dict:
     )
 
 
-def _legacy_density_build_fn(
-    model: str,
-    family_args: dict,
-    extra: dict,
-    input_is_theta: bool,
-) -> Callable:
-    """Return a build function for the models that have no config yet.
-
-    ``mnle`` and ``mnpe`` are reached through the density factories by the
-    deprecated string path of MNLE and MNPE, and are built from flat kwargs by
-    ``build_mnle`` / ``build_mnpe`` rather than from a config.
-    """
-    config = ConditionalFlowConfig.from_kwargs(
-        **{_BUILD_KWARG_ALIASES.get(k, k): v for k, v in family_args.items()},
-        **extra,
-    )
-    builder_kwargs = config.to_dict()
+def _unknown_density_build_fn(model: str) -> Callable:
+    """Preserve the factories' build-time error for unknown model names."""
 
     def build_fn(batch_theta, batch_x):
-        if model not in model_builders:
-            raise NotImplementedError(f"Model {model} is not implemented")
-
-        modeled, condition = (
-            (batch_theta, batch_x) if input_is_theta else (batch_x, batch_theta)
-        )
-        return model_builders[model](
-            batch_x=modeled, batch_y=condition, **builder_kwargs
-        )
+        raise NotImplementedError(f"Model {model} is not implemented")
 
     return build_fn
 
@@ -299,19 +275,27 @@ def likelihood_nn(
         embedding_net=embedding_net,
         num_components=num_components,
     )
-    if model not in _DENSITY_CONFIGS:
-        return _legacy_density_build_fn(
-            model, family_args, kwargs, input_is_theta=False
+    if model in ("mnle", "mnpe"):
+        config = _mixed_config_from_factory_kwargs(
+            family_args=family_args,
+            factory_defaults=_factory_defaults(
+                likelihood_nn, _LIKELIHOOD_FACTORY_FIELDS
+            ),
+            extra=kwargs,
         )
-
-    config = _config_from_factory_kwargs(
-        model,
-        _DENSITY_CONFIGS,
-        "density",
-        family_args=family_args,
-        factory_defaults=_factory_defaults(likelihood_nn, _LIKELIHOOD_FACTORY_FIELDS),
-        extra=kwargs,
-    )
+    elif model in _DENSITY_CONFIGS:
+        config = _config_from_factory_kwargs(
+            model,
+            _DENSITY_CONFIGS,
+            "density",
+            family_args=family_args,
+            factory_defaults=_factory_defaults(
+                likelihood_nn, _LIKELIHOOD_FACTORY_FIELDS
+            ),
+            extra=kwargs,
+        )
+    else:
+        return _unknown_density_build_fn(model)
 
     def build_fn(batch_theta, batch_x):
         # NLE models p(x|theta), so the modeled variable is x.
@@ -411,17 +395,23 @@ def build_fn_snpe_a(batch_theta, batch_x, num_components):
 
         return build_fn_snpe_a
 
-    if model not in _DENSITY_CONFIGS:
-        return _legacy_density_build_fn(model, family_args, kwargs, input_is_theta=True)
-
-    config = _config_from_factory_kwargs(
-        model,
-        _DENSITY_CONFIGS,
-        "density",
-        family_args=family_args,
-        factory_defaults=_factory_defaults(posterior_nn, _POSTERIOR_FACTORY_FIELDS),
-        extra=kwargs,
-    )
+    if model in ("mnle", "mnpe"):
+        config = _mixed_config_from_factory_kwargs(
+            family_args=family_args,
+            factory_defaults=_factory_defaults(posterior_nn, _POSTERIOR_FACTORY_FIELDS),
+            extra=kwargs,
+        )
+    elif model in _DENSITY_CONFIGS:
+        config = _config_from_factory_kwargs(
+            model,
+            _DENSITY_CONFIGS,
+            "density",
+            family_args=family_args,
+            factory_defaults=_factory_defaults(posterior_nn, _POSTERIOR_FACTORY_FIELDS),
+            extra=kwargs,
+        )
+    else:
+        return _unknown_density_build_fn(model)
 
     def build_fn(batch_theta, batch_x):
         # NPE models p(theta|x), so the modeled variable is theta.

sbi/neural_nets/net_builders/estimator_configs.py

@@ -1535,16 +1535,7 @@ def build(
         return _build_mixed_density_estimator(
             batch_x=batch_input,
             batch_y=batch_condition,
-            continuous_config=self.continuous,
-            z_score_y=self.z_score_condition,
-            num_categories_per_variable=self.num_categories_per_variable,
-            embedding_net=self.embedding_net,
-            combined_embedding_net=self.combined_embedding_net,
-            log_transform_x=self.log_transform_x,
-            discrete_hidden_features=self.discrete_hidden_features,
-            discrete_hidden_layers=self.discrete_hidden_layers,
-            combined_embedding_features=self.combined_embedding_features,
-            dropout_probability=self.dropout_probability,
+            config=self,
         )
 
 
@@ -1637,3 +1628,85 @@ def _config_from_factory_kwargs(
         )
 
     return config_cls(**accepted, extra_kwargs=unknown)
+
+
+def _mixed_config_from_factory_kwargs(
+    family_args: dict,
+    factory_defaults: dict,
+    extra: dict,
+) -> MixedConfig:
+    """Build a mixed config from the deprecated factories' flat arguments."""
+    # The flat API typed these as `Optional[int] = None`, so an explicit None
+    # meant "unset" and has to keep behaving like omitting the argument.
+    unset_if_none = (
+        "continuous_hidden_features",
+        "discrete_hidden_features",
+        "combined_embedding_features",
+    )
+    extra = {
+        name: value
+        for name, value in extra.items()
+        if value is not None or name not in unset_if_none
+    }
+    flow_model = extra.pop("flow_model", "nsf")
+
+    mixed_kwargs = {
+        "z_score_condition": family_args["z_score_condition"],
+        "embedding_net": family_args["embedding_net"],
+    }
+    mixed_fields = (
+        "num_categories_per_variable",
+        "combined_embedding_net",
+        "log_transform_x",
+        "discrete_hidden_features",
+        "discrete_hidden_layers",
+        "combined_embedding_features",
+        "dropout_probability",
+    )
+    for name in mixed_fields:
+        if name in extra:
+            mixed_kwargs[name] = extra.pop(name)
+
+    continuous_args = {
+        name: family_args[name]
+        for name in (
+            "z_score_input",
+            "hidden_features",
+            "num_transforms",
+            "num_bins",
+            "num_components",
+        )
+    }
+    continuous_defaults = {name: factory_defaults[name] for name in continuous_args}
+
+    if "continuous_hidden_features" in extra:
+        # The flat API let this override only the continuous width while the
+        # categorical width kept falling back to `hidden_features`.
+        mixed_kwargs.setdefault(
+            "discrete_hidden_features", continuous_args["hidden_features"]
+        )
+        continuous_args["hidden_features"] = extra.pop("continuous_hidden_features")
+
+    dropout_probability = mixed_kwargs.get("dropout_probability")
+    config_cls = _DENSITY_CONFIGS.get(flow_model)
+    if (
+        dropout_probability is not None
+        and config_cls is not None
+        and "dropout_probability" in {f.name for f in fields(config_cls)}
+    ):
+        extra["dropout_probability"] = dropout_probability
+
+    # The flat mixed API passed `tail_bound` to every builder, so the models
+    # that read it saw 10.0 rather than their own narrower default.
+    if config_cls is not None and "tail_bound" in {f.name for f in fields(config_cls)}:
+        extra.setdefault("tail_bound", 10.0)
+
+    continuous = _config_from_factory_kwargs(
+        flow_model,
+        _DENSITY_CONFIGS,
+        "mixed continuous density",
+        family_args=continuous_args,
+        factory_defaults=continuous_defaults,
+        extra=extra,
+    )
+    return MixedConfig(continuous=continuous, **mixed_kwargs)