Merge pull request #2664 from sdaulton/hadamard

Supporting passing full train inputs, rather than only supporting task indices to in `HadamardGaussianLikelihood`

Author: sdaulton

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

@@ -45,6 +45,7 @@ class HadamardGaussianLikelihood(_GaussianLikelihoodBase):
             type of prior.
         noise_constraint: Constraint on the noise value.
         batch_shape: The batch shape of the learned noise parameter (default: []).
+        task_feature_index: The index of the task feature in the input data (default: None).
     """
 
     def __init__(
@@ -53,6 +54,7 @@ def __init__(
         noise_prior: Optional[Prior] = None,
         noise_constraint: Optional[Interval] = None,
         batch_shape: torch.Size = torch.Size(),
+        task_feature_index: Optional[int] = None,
         **kwargs,
     ):
         noise_covar = MultitaskHomoskedasticNoise(
@@ -62,6 +64,7 @@ def __init__(
             batch_shape=batch_shape,
         )
         self.num_tasks = num_tasks
+        self.task_feature_index = task_feature_index
         super().__init__(noise_covar=noise_covar, **kwargs)
 
     @property
@@ -81,17 +84,26 @@ def raw_noise(self, value: torch.Tensor) -> None:
         self.noise_covar.initialize(raw_noise=value)
 
     def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any) -> LinearOperator:
-        # params contains task indexes, shape (*task_batch_shape, num_data, 1)
-        if not params or not params[0]:
+        # params contains input data, shape (*task_batch_shape, num_data, d)
+        if len(params) == 0 or len(params[0]) == 0:
             raise ValueError("Task indices must be provided.")
-        task_idcs = params[0][-1]
+        task_idcs = params[0] if torch.is_tensor(params[0]) else params[0][0]
+        if task_idcs.shape[-1] > 1:
+            if self.task_feature_index is None:
+                raise ValueError("Task indices must be a single dimension if task_feature_index is not provided.")
+            # handle case where full inputs are passed in,
+            # rather than just the task indices
+            task_idcs = task_idcs[..., self.task_feature_index].unsqueeze(-1)
+        task_idcs = task_idcs.long()
         _check_task_indices(task_idcs, base_shape)
-
-        # squeeze to remove the `1` dimension returned by `MultitaskHomoskedasticNoise`
-        noise_base_covar_matrix = self.noise_covar(*params, shape=base_shape, **kwargs).squeeze(
-            -4
-        )  # (*batch_shape, num_tasks, num_data, num_data)
-        all_tasks = torch.arange(self.num_tasks).unsqueeze(-1)  # (num_tasks, 1)
+        noise_base_covar_matrix = self.noise_covar(task_idcs, *params[1:], shape=base_shape, **kwargs)
+        if self.num_tasks > 1:
+            # squeeze to remove the `1` dimension returned by `MultitaskHomoskedasticNoise`
+            # when there is more than 1 task
+            noise_base_covar_matrix = noise_base_covar_matrix.squeeze(
+                -4
+            )  # (*batch_shape, num_tasks, num_data, num_data)
+        all_tasks = torch.arange(self.num_tasks, device=task_idcs.device).unsqueeze(-1)  # (num_tasks, 1)
         diag = torch.eq(all_tasks, task_idcs.mT)  # (num_tasks, num_data)
         mask = DiagLinearOperator(diag)  # (num_tasks, num_data, num_data)
         return (noise_base_covar_matrix @ mask).sum(dim=-3)

test/likelihoods/test_hadamard_multitask_gaussian_likelihood.py

@@ -111,8 +111,11 @@ def test_marginal_variance(self):
         with self.assertRaises(ValueError):
             likelihood(input, [])
 
-        with self.assertRaises(ValueError):
-            likelihood(input, [task_idcs.float()])
-
         with self.assertRaises(ValueError):
             likelihood(input, [task_idcs.squeeze(-1)])
+
+        # test with task feature and full input
+        likelihood = HadamardGaussianLikelihood(num_tasks=4, task_feature_index=1)
+        likelihood.noise = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
+        X = torch.cat([torch.zeros_like(task_idcs), task_idcs], dim=-1)
+        self.assertEqual(variance, likelihood(input, [X]).variance)