Extend UIGP to multiple uncertain parameters (with different input variances)

gpax/models/uigp.py

@@ -110,13 +110,23 @@ def model(self, X: jnp.ndarray, y: jnp.ndarray = None, **kwargs: float) -> None:
             obs=y,
         )
 
-    def _sample_x(self, X):
+    def _sample_x(self, X: jnp.ndarray) -> jnp.ndarray:
+        """
+        Samples new input values (X_prime) based on the original inputs (X)
+        and prior belief about the uncertainty in those inputs.
+        """
+        n_samples, n_features = X.shape
         if self.sigma_x_prior_dist is not None:
             sigma_x_dist = self.sigma_x_prior_dist
         else:
-            sigma_x_dist = dist.HalfNormal(.1)
-        sigma_x = numpyro.sample("sigma_x", sigma_x_dist)
-        return numpyro.sample("X_prime", dist.Normal(X, sigma_x))
+            sigma_x_dist = dist.HalfNormal(.1 * jnp.ones(n_features))
+        # Sample variances independently for each feature dimension
+        with numpyro.plate("feature_variance_plate", self.kernel_dim):
+            sigma_x = numpyro.sample("sigma_x", sigma_x_dist)
+            # Sample input data using the sampled variances
+            with numpyro.plate("X_prime_plate", n_samples, dim=-2):
+                X_prime = numpyro.sample("X_prime", dist.Normal(X, sigma_x))
+        return X_prime
 
     def get_mvn_posterior(
         self, X_new: jnp.ndarray, params: Dict[str, jnp.ndarray], noiseless: bool = False, **kwargs: float

tests/test_uigp.py

@@ -20,20 +20,30 @@ def get_dummy_data():
     return jnp.array(X_prime), jnp.array(y)
 
 
+@pytest.mark.parametrize("n_features", [1, 5])
+def test_sample_x(n_features):
+    X = onp.random.randn(32, n_features)
+    m = UIGP(n_features, 'RBF')
+    with numpyro.handlers.seed(rng_seed=0):
+        X_prime = m._sample_x(X)
+    assert_(isinstance(X_prime, jnp.ndarray))
+    assert_(X_prime.shape[-1], n_features)
+
+
 def test_fit():
     rng_key = get_keys()[0]
     X, y = get_dummy_data()
     m = UIGP(1, 'RBF')
     m.fit(rng_key, X, y, num_warmup=10, num_samples=10)
-    assert m.mcmc is not None
+    assert_(m.mcmc is not None)
 
 
 def test_fit_with_custom_sigma_x_prior():
     rng_key = get_keys()[0]
     X, y = get_dummy_data()
     m = UIGP(1, 'RBF', sigma_x_prior_dist=dist.HalfNormal(0.55))
     m.fit(rng_key, X, y, num_warmup=10, num_samples=10)
-    assert m.mcmc is not None
+    assert_(m.mcmc is not None)
 
 
 def test_get_mvn_posterior():
@@ -52,8 +62,8 @@ def test_get_mvn_posterior():
     m.X_train = X
     m.y_train = y
     mean, cov = m.get_mvn_posterior(X_test, params)
-    assert isinstance(mean, jnp.ndarray)
-    assert isinstance(cov, jnp.ndarray)
+    assert_(isinstance(mean, jnp.ndarray))
+    assert_(isinstance(cov, jnp.ndarray))
     assert_equal(mean.shape, (X_test.shape[0],))
     assert_equal(cov.shape, (X_test.shape[0], X_test.shape[0]))
 
@@ -75,7 +85,7 @@ def test_predict_single_sample(noiseless):
     m.X_train = X
     m.y_train = y
     mean, sample = m._predict(key, X_test, params, 5, noiseless)
-    assert isinstance(mean, jnp.ndarray)
-    assert isinstance(sample, jnp.ndarray)
+    assert_(isinstance(mean, jnp.ndarray))
+    assert_(isinstance(sample, jnp.ndarray))
     assert_equal(mean.shape, (X_test.shape[0],))
     assert_equal(sample.shape, (5, X_test.shape[0]))