fix shape issue

Author: radka-j

autoemulate/core/metrics.py

@@ -325,7 +325,7 @@ def __call__(
         y_true: TensorLike
             True target values.
         metric_params: MetricParams
-            Metric parameters including: y_train.
+            Metric parameters including: y_train and reduction.
 
         Returns
         -------
@@ -352,24 +352,18 @@ def __call__(
         y_true = y_true.unsqueeze(-1) if y_true.ndim == 1 else y_true
 
         # Compute mean log loss
-        mean_log_loss = y_pred.log_prob(y_true).mean(dim=0)
+        mean_log_loss = y_pred.log_prob(y_true).mean()
 
         # If no training data, return mean log loss
         if metric_params.y_train is None:
             return mean_log_loss
 
         # Ensure 2D y_train for consistent handling
-        y_train = (
-            metric_params.y_train.unsqueeze(-1)
-            if metric_params.y_train.ndim == 1
-            else metric_params.y_train
-        )
-
-        y_train_mean = y_train.mean()
+        y_train_mean = metric_params.y_train.mean(dim=0, keepdim=True).view(1, -1)
 
         # following GPyTorch implementation, use global variance rather than per task
         # https://github.com/cornellius-gp/gpytorch/blob/c0fb6c64311fdbef2862fd3ba2bd613fbd081e79/gpytorch/metrics/metrics.py#L60
-        y_train_var = y_train.var()
+        y_train_var = metric_params.y_train.var()
 
         # Avoid numerical issues
         y_train_var = torch.clamp(y_train_var, min=1e-6)

tests/core/test_metrics.py

@@ -557,14 +557,9 @@ def test_msll_with_1d_inputs():
     )
 
     msll = MSLL(
-        y_pred, y_true, metric_params=MetricParams(y_train=y_train, reduction="none")
-    )
-
-    assert isinstance(msll, torch.Tensor)
-    assert msll.shape == torch.Size([1])
-
-    msll = MSLL(
-        y_pred, y_true, metric_params=MetricParams(y_train=y_train, reduction="mean")
+        y_pred,
+        y_true,
+        metric_params=MetricParams(y_train=y_train),
     )
 
     assert isinstance(msll, torch.Tensor)
@@ -579,14 +574,9 @@ def test_msll_with_2d_inputs():
     y_pred = Normal(loc=y_true, scale=torch.ones_like(y_true))
 
     msll = MSLL(
-        y_pred, y_true, metric_params=MetricParams(y_train=y_train, reduction="none")
-    )
-
-    assert isinstance(msll, torch.Tensor)
-    assert msll.shape == torch.Size([n_outputs])
-
-    msll = MSLL(
-        y_pred, y_true, metric_params=MetricParams(y_train=y_train, reduction="mean")
+        y_pred,
+        y_true,
+        metric_params=MetricParams(y_train=y_train),
     )
 
     assert isinstance(msll, torch.Tensor)
@@ -607,7 +597,7 @@ def test_msll_raises_for_non_distribution():
 def test_msll_perfect_prediction_with_training_data():
     """Test MSLL when predictions perfectly match true values."""
     y_true = torch.tensor([1.0, 2.0, 3.0]).view(-1, 1)
-    y_train = torch.tensor([0.5, 1.5, 2.5]).view(-1, 1)
+    y_train = torch.tensor([0.5, 1.5, 2.5])
 
     # Perfect prediction: distribution centered at true values with small variance
     y_pred = Normal(loc=y_true, scale=torch.ones_like(y_true) * 0.01)
@@ -621,7 +611,7 @@ def test_msll_perfect_prediction_with_training_data():
 def test_msll_poor_prediction_with_training_data():
     """Test MSLL when predictions are far from true values."""
     y_true = torch.tensor([1.0, 2.0, 3.0]).view(-1, 1)
-    y_train = torch.tensor([1.0, 2.0, 3.0]).view(-1, 1)
+    y_train = torch.tensor([1.0, 2.0, 3.0])
 
     # Poor prediction: distribution centered far from true values
     y_pred = Normal(