Fix pickle error in TorchPosterior (meta-pytorch#1644)

Summary:
## Motivation

Fixes meta-pytorch#1639. The issue was that pickle uses `__getstate__` and `__setstate__` to save / load the objects. In `TorchPosterior`, `__getattr__` was taking precedence, so instead of calling `TorchPosterior.__set/getstate__` it was calling `self.distribution.__set/getstate__`. While loading the posterior, we start with a `TorchPosterior` that does not have any attributes --including `distribution`--, so calling `self.distribution` was leading to an infinite loop of calling `TorchPosterior` with `name=distribution` (since it calls `getattr(self.distribution, name)`).

Pull Request resolved: meta-pytorch#1644

Test Plan: Units.

Reviewed By: Balandat

Differential Revision: D42732823

Pulled By: saitcakmak

fbshipit-source-id: 20183a7d5368a1e92dc36cd9726d0caf33f791d5

Author: saitcakmak

botorch/posteriors/torch.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import torch
 from botorch.posteriors.posterior import Posterior
@@ -78,6 +78,20 @@ def __getattr__(self, name: str) -> Any:
         """
         return getattr(self.distribution, name)
 
+    def __getstate__(self) -> Dict[str, Any]:
+        r"""A minimal utility to support pickle protocol.
+
+        Pickle uses `__get/setstate__` to serialize / deserialize the objects.
+        Since we define `__getattr__` above, it takes precedence over these
+        methods, and we end up in an infinite loop unless we also define
+        `__getstate__` and `__setstate__`.
+        """
+        return self.__dict__
+
+    def __setstate__(self, d: Dict[str, Any]) -> None:
+        r"""A minimal utility to support pickle protocol."""
+        self.__dict__ = d
+
     def quantile(self, value: Tensor) -> Tensor:
         r"""Compute quantiles of the distribution.
 

test/posteriors/test_torch_posterior.py

@@ -4,6 +4,9 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+import tempfile
+import unittest
 
 import torch
 from botorch.posteriors.torch import TorchPosterior
@@ -55,3 +58,20 @@ def test_torch_posterior(self):
                 [posterior.distribution.log_prob(q).exp() for q in q_value], dim=0
             )
             self.assertAllClose(posterior.density(q_value), expected)
+
+    @unittest.skipIf(os.name == "nt", "Pickle test is not supported on Windows.")
+    def test_pickle(self) -> None:
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"dtype": dtype, "device": self.device}
+            posterior = TorchPosterior(Exponential(rate=torch.rand(1, 2, **tkwargs)))
+            with tempfile.NamedTemporaryFile() as tmp_file:
+                torch.save(posterior, tmp_file.name)
+                loaded_posterior = torch.load(tmp_file.name)
+            self.assertEqual(posterior.dtype, loaded_posterior.dtype)
+            self.assertEqual(posterior.device, loaded_posterior.device)
+            self.assertTrue(
+                torch.equal(
+                    posterior.distribution.rate,
+                    loaded_posterior.distribution.rate,
+                )
+            )