Merge pull request #1992 from cornellius-gp/pickle_constant_mean

Make it possible to pickle constant mean with prior

Author: Balandat

gpytorch/means/constant_mean.py

@@ -12,7 +12,13 @@ def __init__(self, prior=None, batch_shape=torch.Size(), **kwargs):
         self.batch_shape = batch_shape
         self.register_parameter(name="constant", parameter=torch.nn.Parameter(torch.zeros(*batch_shape, 1)))
         if prior is not None:
-            self.register_prior("mean_prior", prior, "constant")
+            self.register_prior("mean_prior", prior, self._constant_param, self._constant_closure)
+
+    def _constant_param(self, m):
+        return m.constant
+
+    def _constant_closure(self, m, value):
+        return m.constant.data.fill_(value)
 
     def forward(self, input):
         if input.shape[:-2] == self.batch_shape:

test/means/test_constant_mean.py

@@ -1,17 +1,27 @@
 #!/usr/bin/env python3
 
+import pickle
 import unittest
 
 import torch
 
 from gpytorch.means import ConstantMean
+from gpytorch.priors import NormalPrior
 from gpytorch.test.base_mean_test_case import BaseMeanTestCase
 
 
 class TestConstantMean(BaseMeanTestCase, unittest.TestCase):
     def create_mean(self):
         return ConstantMean()
 
+    def test_prior(self):
+        prior = NormalPrior(0.0, 1.0)
+        mean = ConstantMean(prior=prior)
+        self.assertEqual(mean.mean_prior, prior)
+        pickle.loads(pickle.dumps(mean))  # Should be able to pickle and unpickle with a prior
+        mean._constant_closure(mean, 1.234)
+        self.assertAlmostEqual(mean.constant.item(), 1.234)
+
 
 class TestConstantMeanBatch(BaseMeanTestCase, unittest.TestCase):
     batch_shape = torch.Size([3])