Support multi-output y in obs_variance computation

Expand leverage shape (n,) -> (n, 1) for broadcasting with (n, p)
residuals in all three conditional classes. The triangular solves
and _process_sigma already handle matrix RHS naturally.

Author: katosh

mellon/conditional.py

@@ -237,6 +237,8 @@ def _compute_obs_variance(self, x, y, mu, cov_func, sigma, jitter, weights):
 
         # Corrected squared residuals (HC3 estimator)
         residual = y - prediction
+        if residual.ndim > h.ndim:
+            h = h[..., None]
         corrected_r2 = residual**2 / (1 - h) ** 2
 
         # Fit second GP to corrected_r2 with noise regularization sigma.
@@ -456,6 +458,8 @@ def _compute_obs_variance(self, x, y, xu, mu, cov_func, sigma, jitter, weights):
 
         # Corrected squared residuals (HC3 estimator)
         residual = y - prediction
+        if residual.ndim > h.ndim:
+            h = h[..., None]
         corrected_r2 = residual**2 / (1 - h) ** 2
 
         # Fit second GP on landmarks to corrected_r2 with noise sigma.
@@ -675,6 +679,8 @@ def _compute_obs_variance(self, x, y, xu, mu, cov_func, sigma, jitter, weights):
 
         # Corrected squared residuals (HC3 estimator)
         residual = y - prediction
+        if residual.ndim > h.ndim:
+            h = h[..., None]
         corrected_r2 = residual**2 / (1 - h) ** 2
 
         # Fit second GP on landmarks to corrected_r2 with noise sigma.

tests/test_leverage.py

@@ -247,6 +247,27 @@ def test_estimator_get_obs_variance(setup_data):
     )
 
 
+def test_obs_variance_multi_output():
+    """obs_variance should work with multi-output y of shape (n, p)."""
+    n, d, p = 80, 3, 4
+    key = jax.random.PRNGKey(42)
+    k1, k2 = jax.random.split(key)
+    X = jax.random.normal(k1, (n, d))
+    y = jax.random.normal(k2, (n, p))
+
+    # Full GP
+    est = mellon.FunctionEstimator(sigma=1.0, n_landmarks=0, obs_variance=True)
+    est.fit(X, y)
+    var = est.predict.obs_variance(X)
+    assert var.shape == (n, p), f"Full GP: expected ({n}, {p}), got {var.shape}"
+
+    # Sparse GP
+    est_s = mellon.FunctionEstimator(sigma=1.0, n_landmarks=20, obs_variance=True)
+    est_s.fit(X, y)
+    var_s = est_s.predict.obs_variance(X)
+    assert var_s.shape == (n, p), f"Sparse GP: expected ({n}, {p}), got {var_s.shape}"
+
+
 def test_fit_obs_variance_override(setup_data):
     """fit(x, y, obs_variance=True) should override constructor default."""
     X, y = setup_data