Add fixed noise models to reverse registry (#1938)

Summary:
Pull Request resolved: #1938

The changes in D50431137 removed these models from the values of `MODEL_REGISTRY`, which lead to them getting removed from the reverse registry as well. Since these models appear in the DB for the past experiments, we need to manually add them to the reverse registry.

Reviewed By: sdaulton

Differential Revision: D50742593

fbshipit-source-id: 67cb966ae2d78c9234dbdc9f8caa62cf26335a7f

ax/storage/botorch_modular_registry.py

@@ -195,7 +195,10 @@
 
 
 REVERSE_MODEL_REGISTRY: Dict[str, Type[Model]] = {
-    v: k for k, v in MODEL_REGISTRY.items()
+    "FixedNoiseGP": SingleTaskGP,
+    "FixedNoiseMultiFidelityGP": SingleTaskMultiFidelityGP,
+    "FixedNoiseMultiTaskGP": MultiTaskGP,
+    **{v: k for k, v in MODEL_REGISTRY.items()},
 }
 
 

ax/utils/testing/modeling_stubs.py

@@ -361,16 +361,18 @@ def get_surrogate_as_dict() -> Dict[str, Any]:
     }
 
 
-def get_surrogate_spec_as_dict() -> Dict[str, Any]:
+def get_surrogate_spec_as_dict(model_class: Optional[str] = None) -> Dict[str, Any]:
     """
     For use ensuring backwards compatibility when loading SurrogateSpec
     with input_transform and outcome_transform kwargs.
     """
+    if model_class is None:
+        model_class = "SingleTaskGP"
     return {
         "__type": "SurrogateSpec",
         "botorch_model_class": {
             "__type": "Type[Model]",
-            "index": "SingleTaskGP",
+            "index": model_class,
             "class": "<class 'botorch.models.model.Model'>",
         },
         "botorch_model_kwargs": {},