Storage support for RBFKernel & LogNormalPrior (facebook#2616)

Summary:
Pull Request resolved: facebook#2616

Adds storage support for these components that'll soon be used by default. The storage support is only relevant in cases where they are manually specified.

I also removed `REVERSE_GPYTORCH_COMPONENT_REGISTRY` since I did not find it to affect anything. We need `GPYTORCH_COMPONENT_REGISTRY` but the reverse does not seem to do much.

Reviewed By: mpolson64

Differential Revision: D60458209

fbshipit-source-id: a8f8cd826656d497a35d60fa54e3fc901b88481d

Author: saitcakmak

ax/storage/botorch_modular_registry.py

@@ -79,6 +79,7 @@
 # Miscellaneous BoTorch imports
 from gpytorch.constraints import Interval
 from gpytorch.kernels.kernel import Kernel
+from gpytorch.kernels.rbf_kernel import RBFKernel
 from gpytorch.likelihoods.gaussian_likelihood import GaussianLikelihood
 from gpytorch.likelihoods.likelihood import Likelihood
 
@@ -156,6 +157,7 @@
 
 KERNEL_REGISTRY: Dict[Type[Kernel], str] = {
     ScaleMaternKernel: "ScaleMaternKernel",
+    RBFKernel: "RBFKernel",
 }
 
 LIKELIHOOD_REGISTRY: Dict[Type[GaussianLikelihood], str] = {
@@ -200,6 +202,7 @@
     Model: MODEL_REGISTRY,
     Interval: GPYTORCH_COMPONENT_REGISTRY,
     GammaPrior: GPYTORCH_COMPONENT_REGISTRY,
+    LogNormalPrior: GPYTORCH_COMPONENT_REGISTRY,
     InputTransform: INPUT_TRANSFORM_REGISTRY,
     OutcomeTransform: OUTCOME_TRANSFORM_REGISTRY,
 }
@@ -239,12 +242,6 @@
     v: k for k, v in LIKELIHOOD_REGISTRY.items()
 }
 
-
-REVERSE_GPYTORCH_COMPONENT_REGISTRY: Dict[str, Type[torch.nn.Module]] = {
-    v: k for k, v in GPYTORCH_COMPONENT_REGISTRY.items()
-}
-
-
 REVERSE_INPUT_TRANSFORM_REGISTRY: Dict[str, Type[InputTransform]] = {
     v: k for k, v in INPUT_TRANSFORM_REGISTRY.items()
 }
@@ -264,8 +261,6 @@
     Likelihood: REVERSE_LIKELIHOOD_REGISTRY,
     MarginalLogLikelihood: REVERSE_MLL_REGISTRY,
     Model: REVERSE_MODEL_REGISTRY,
-    Interval: REVERSE_GPYTORCH_COMPONENT_REGISTRY,
-    GammaPrior: REVERSE_GPYTORCH_COMPONENT_REGISTRY,
     InputTransform: REVERSE_INPUT_TRANSFORM_REGISTRY,
     OutcomeTransform: REVERSE_OUTCOME_TRANSFORM_REGISTRY,
 }

ax/storage/json_store/decoders.py

@@ -41,6 +41,7 @@
 from botorch.models.transforms.input import ChainedInputTransform, InputTransform
 from botorch.models.transforms.outcome import ChainedOutcomeTransform, OutcomeTransform
 from botorch.utils.types import _DefaultType, DEFAULT
+from torch.distributions.transformed_distribution import TransformedDistribution
 
 logger: logging.Logger = get_logger(__name__)
 
@@ -279,6 +280,11 @@ def botorch_component_from_json(botorch_class: Any, json: Dict[str, Any]) -> Typ
                 for k, v in state_dict.items()
             }
         )
+    if issubclass(botorch_class, TransformedDistribution):
+        # Extract the transformed attributes for transformed priors.
+        for k in list(state_dict.keys()):
+            if k.startswith("_transformed_"):
+                state_dict[k[13:]] = state_dict.pop(k)
     class_path = json.pop("class")
     init_args = inspect.signature(botorch_class).parameters
     required_args = {

ax/storage/json_store/registry.py

@@ -171,7 +171,7 @@
 from gpytorch.constraints import Interval
 from gpytorch.likelihoods.likelihood import Likelihood
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
-from gpytorch.priors.torch_priors import GammaPrior
+from gpytorch.priors.torch_priors import GammaPrior, LogNormalPrior
 
 
 # pyre-fixme[5]: Global annotation cannot contain `Any`.
@@ -206,6 +206,7 @@
     Interval: botorch_component_to_dict,
     JenattonMetric: metric_to_dict,
     L2NormMetric: metric_to_dict,
+    LogNormalPrior: botorch_component_to_dict,
     MapData: map_data_to_dict,
     MapKeyInfo: map_key_info_to_dict,
     MapMetric: metric_to_dict,
@@ -327,6 +328,7 @@
     "LifecycleStage": LifecycleStage,
     "ListSurrogate": Surrogate,  # For backwards compatibility
     "L2NormMetric": L2NormMetric,
+    "LogNormalPrior": LogNormalPrior,
     "MapData": MapData,
     "MapMetric": MapMetric,
     "MapKeyInfo": MapKeyInfo,

ax/storage/json_store/tests/test_json_store.py

@@ -114,6 +114,7 @@
     get_sum_constraint2,
     get_surrogate,
     get_surrogate_spec_with_default,
+    get_surrogate_spec_with_lognormal,
     get_synthetic_runner,
     get_threshold_early_stopping_strategy,
     get_trial,
@@ -238,6 +239,7 @@
 class JSONStoreTest(TestCase):
     def setUp(self) -> None:
         super().setUp()
+        self.maxDiff = None
         self.experiment = get_experiment_with_batch_and_single_trial()
 
     def test_JSONEncodeFailure(self) -> None:
@@ -526,32 +528,35 @@ def test_EncodeDecodeSet(self) -> None:
     def test_encode_decode_surrogate_spec(self) -> None:
         # Test SurrogateSpec separately since the GPyTorch components
         # fail simple equality checks.
-        org_object = get_surrogate_spec_with_default()
-        converted_object = object_from_json(
-            object_to_json(
-                org_object,
-                encoder_registry=CORE_ENCODER_REGISTRY,
-                class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
-            ),
-            decoder_registry=CORE_DECODER_REGISTRY,
-            class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
-        )
-        org_as_dict = dataclasses.asdict(org_object)
-        converted_as_dict = dataclasses.asdict(converted_object)
-        # Covar module kwargs will fail comparison. Manually compare.
-        org_covar_kwargs = org_as_dict.pop("covar_module_kwargs")
-        converted_covar_kwargs = converted_as_dict.pop("covar_module_kwargs")
-        self.assertEqual(org_covar_kwargs.keys(), converted_covar_kwargs.keys())
-        for k in org_covar_kwargs:
-            org_ = org_covar_kwargs[k]
-            converted_ = converted_covar_kwargs[k]
-            if isinstance(org_, torch.nn.Module):
-                self.assertEqual(org_.__class__, converted_.__class__)
-                self.assertEqual(org_.__dict__, converted_.__dict__)
-            else:
-                self.assertEqual(org_, converted_)
-        # Compare the rest.
-        self.assertEqual(org_as_dict, converted_as_dict)
+        for org_object in (
+            get_surrogate_spec_with_default(),
+            get_surrogate_spec_with_lognormal(),
+        ):
+            converted_object = object_from_json(
+                object_to_json(
+                    org_object,
+                    encoder_registry=CORE_ENCODER_REGISTRY,
+                    class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+                ),
+                decoder_registry=CORE_DECODER_REGISTRY,
+                class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+            )
+            org_as_dict = dataclasses.asdict(org_object)
+            converted_as_dict = dataclasses.asdict(converted_object)
+            # Covar module kwargs will fail comparison. Manually compare.
+            org_covar_kwargs = org_as_dict.pop("covar_module_kwargs")
+            converted_covar_kwargs = converted_as_dict.pop("covar_module_kwargs")
+            self.assertEqual(org_covar_kwargs.keys(), converted_covar_kwargs.keys())
+            for k in org_covar_kwargs:
+                org_ = org_covar_kwargs[k]
+                converted_ = converted_covar_kwargs[k]
+                if isinstance(org_, torch.nn.Module):
+                    self.assertEqual(org_.__class__, converted_.__class__)
+                    self.assertEqual(org_.state_dict(), converted_.state_dict())
+                else:
+                    self.assertEqual(org_, converted_)
+            # Compare the rest.
+            self.assertEqual(org_as_dict, converted_as_dict)
 
     def test_RegistryAdditions(self) -> None:
         class MyRunner(Runner):

ax/utils/testing/core_stubs.py

@@ -122,9 +122,10 @@
 from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.types import DEFAULT
 from gpytorch.constraints import Interval
+from gpytorch.kernels.rbf_kernel import RBFKernel
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
-from gpytorch.priors.torch_priors import GammaPrior
+from gpytorch.priors.torch_priors import GammaPrior, LogNormalPrior
 
 logger: Logger = get_logger(__name__)
 
@@ -2251,6 +2252,18 @@ def get_surrogate_spec_with_default() -> SurrogateSpec:
     )
 
 
+def get_surrogate_spec_with_lognormal() -> SurrogateSpec:
+    return SurrogateSpec(
+        botorch_model_class=SingleTaskGP,
+        covar_module_class=RBFKernel,
+        covar_module_kwargs={
+            "ard_num_dims": DEFAULT,
+            "lengthscale_prior": LogNormalPrior(-4.0, 1.0),
+            "batch_shape": DEFAULT,
+        },
+    )
+
+
 def get_acquisition_type() -> Type[Acquisition]:
     return Acquisition