FixedNoiseGaussianLikelihood Noise Not Handled Correctly for LikelihoodLists

[CalumHarvey]

🐛 Bug

When calling FixedNoiseGaussianLikelihood with a LikelihoodList, the noise parameter is passed incorrectly resulting in an attempt to apply index references to a dictionary.

To reproduce

Modification to the example here with a FixedNoiseGaussianLikelihood likelihood.

train_x1 = torch.linspace(0, 0.95, 50) + 0.05 * torch.rand(50)
train_x2 = torch.linspace(0, 0.95, 25) + 0.05 * torch.rand(25)

train_y1 = torch.sin(train_x1 * (2 * math.pi)) + 0.2 * torch.randn_like(train_x1)
train_y2 = torch.cos(train_x2 * (2 * math.pi)) + 0.2 * torch.randn_like(train_x2)

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

noise1 = torch.ones(train_y1.shape) * 0.1
likelihood1 = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(noise=noise1)
model1 = ExactGPModel(train_x1, train_y1, likelihood1)

noise2 = torch.ones(train_y2.shape) * 0.1
likelihood2 = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(noise=noise1)
model2 = ExactGPModel(train_x2, train_y2, likelihood2)

model = gpytorch.models.IndependentModelList(model1, model2)
likelihood = gpytorch.likelihoods.LikelihoodList(model1.likelihood, model2.likelihood)

mll = SumMarginalLogLikelihood(likelihood, model)

model.eval()
likelihood.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    test_x = torch.linspace(0, 1, 51)
    test_noise = torch.ones(2, len(test_x))
    predictions = likelihood(*model(test_x, test_x), noise=test_noise)

** Stack trace/error message **

{
	"name": "KeyError",
	"message": "0",
	"stack": "---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[26], line 42
     40 test_x = torch.linspace(0, 1, 51)
     41 test_noise = torch.ones(2, len(test_x))
---> 42 predictions = likelihood(*model(test_x, test_x), noise=test_noise)

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\likelihood_list.py:52, in LikelihoodList.__call__(self, *args, **kwargs)
     50     noise = kwargs.pop(\"noise\")
     51     # if noise kwarg is passed, assume it's an iterable of noise tensors
---> 52     return [
     53         likelihood(*args_, {**kwargs, \"noise\": noise_})
     54         for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
     55     ]
     56 else:
     57     return [
     58         likelihood(*args_, **kwargs)
     59         for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
     60     ]

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\likelihood_list.py:53, in <listcomp>(.0)
     50     noise = kwargs.pop(\"noise\")
     51     # if noise kwarg is passed, assume it's an iterable of noise tensors
     52     return [
---> 53         likelihood(*args_, {**kwargs, \"noise\": noise_})
     54         for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
     55     ]
     56 else:
     57     return [
     58         likelihood(*args_, **kwargs)
     59         for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
     60     ]

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\likelihood.py:367, in Likelihood.__call__(self, input, *args, **kwargs)
    356 # Marginal
    357 elif any(
    358     [
    359         isinstance(input, MultivariateNormal),
   (...)
    365     ]
    366 ):
--> 367     return self.marginal(input, *args, **kwargs)  # pyre-ignore[6]
    368 # Error
    369 else:
    370     raise RuntimeError(
    371         \"Likelihoods expects a MultivariateNormal or Normal input to make marginal predictions, or a \"
    372         \"torch.Tensor for conditional predictions. Got a {}\".format(input.__class__.__name__)
    373     )

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\gaussian_likelihood.py:359, in FixedNoiseGaussianLikelihood.marginal(self, function_dist, *args, **kwargs)
    355 def marginal(self, function_dist: MultivariateNormal, *args: Any, **kwargs: Any) -> MultivariateNormal:
    356     r\"\"\"
    357     :return: Analytic marginal :math:`p(\\mathbf y)`.
    358     \"\"\"
--> 359     return super().marginal(function_dist, *args, **kwargs)

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\gaussian_likelihood.py:116, in _GaussianLikelihoodBase.marginal(self, function_dist, *params, **kwargs)
    114 def marginal(self, function_dist: MultivariateNormal, *params: Any, **kwargs: Any) -> MultivariateNormal:
    115     mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
--> 116     noise_covar = self._shaped_noise_covar(mean.shape, *params, **kwargs)
    117     full_covar = covar + noise_covar
    118     return function_dist.__class__(mean, full_covar)

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\gaussian_likelihood.py:342, in FixedNoiseGaussianLikelihood._shaped_noise_covar(self, base_shape, *params, **kwargs)
    338 else:
    339     # here shape[:-1] is the batch shape requested, and shape[-1] is `n`, the number of points
    340     shape = base_shape
--> 342 res = self.noise_covar(*params, shape=shape, **kwargs)
    344 if self.second_noise_covar is not None:
    345     res = res + self.second_noise_covar(*params, shape=shape, **kwargs)

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\
oise_models.py:182, in FixedGaussianNoise.__call__(self, shape, *params, **kwargs)
    177 def __call__(
    178     self, *params: Any, shape: Optional[torch.Size] = None, **kwargs: Any
    179 ) -> Union[Tensor, LinearOperator]:
    180     # For corredct typing
    181     # print(params)
--> 182     return super().__call__(*params, shape=shape, **kwargs)

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\module.py:32, in Module.__call__(self, *inputs, **kwargs)
     30 def __call__(self, *inputs, **kwargs) -> Union[Tensor, Distribution, LinearOperator]:
     31     # print(inputs)
---> 32     outputs = self.forward(*inputs, **kwargs)
     33     if isinstance(outputs, list):
     34         return [_validate_module_outputs(output) for output in outputs]

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\
oise_models.py:163, in FixedGaussianNoise.forward(self, shape, noise, *params, **kwargs)
    159 def forward(
    160     self, *params: Any, shape: Optional[torch.Size] = None, noise: Optional[Tensor] = None, **kwargs: Any
    161 ) -> DiagLinearOperator:
    162     if shape is None:
--> 163         p = params[0] if torch.is_tensor(params[0]) else params[0][0]
    164         shape = p.shape if len(p.shape) == 1 else p.shape[:-1]
    166     if noise is not None:

KeyError: 0"
}

Expected Behavior

Expected to make predictions without error.

System information

Please complete the following information:

• GPyTorch Version 1.12
• PyTorch Version 2.4.0
• Windows 10

Additional context

The error can be fixed by adding an addition check to the FixedGaussianNoise.forward method as shown below, however, I am unsure of the unintended effects of making this change.

def forward(
        self, *params: Any, shape: Optional[torch.Size] = None, noise: Optional[Tensor] = None, **kwargs: Any
    ) -> DiagLinearOperator:
        if shape is None:
            if type(params[0]) == dict:
                p = list(params[0].values())[0]
            else:
                p = params[0] if torch.is_tensor(params[0]) else params[0][0]
            shape = p.shape if len(p.shape) == 1 else p.shape[:-1]

        if noise is not None:
            return DiagLinearOperator(noise)
        elif shape[-1] == self.noise.shape[-1]:
            return DiagLinearOperator(self.noise)
        else:
            return ZeroLinearOperator()

[Balandat]

The following code handles this if the noise is passed in as an iterable of noise tensors (one for each of the likelihoods):

gpytorch/gpytorch/likelihoods/likelihood_list.py

Lines 31 to 35 in b017b9c

	# if noise kwarg is passed, assume it's an iterable of noise tensors
	return [
	likelihood.forward(*args_, {**kwargs, "noise": noise_})
	for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
	]

Does this help?

[CalumHarvey]

The same error occurs if the noise is passed as a list of noise tensors which length_safe_zip should be able to iterate over:

    test_noise = torch.ones(len(test_x))
    predictions = likelihood(*model(test_x, test_x), noise=[test_noise, test_noise])

Having looked into it more, there appears to be argument passing issue with the calling of the likelihood function. The likelihood is taking the arguments likelihood(*args_, {**kwargs, \"noise\": noise_}) but the function it is calling has signature Likelihood.__call__(self, input, *args, **kwargs) resulting in *args being passed as input and {**kwargs, \"noise\": noise_} being passed as *args.

Does this seem correct or am I missing something? I have included the relevant portion of the trace below:

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\likelihood_list.py:53, in <listcomp>(.0)
     50     noise = kwargs.pop(\"noise\")
     51     # if noise kwarg is passed, assume it's an iterable of noise tensors
     52     return [
---> 53         likelihood(*args_, {**kwargs, \"noise\": noise_})
     54         for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
     55     ]
     56 else:
     57     return [
     58         likelihood(*args_, **kwargs)
     59         for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
     60     ]

File c:\\Users\\ch982\\.conda\\envs\\EDFA_env\\Lib\\site-packages\\gpytorch\\likelihoods\\likelihood.py:367, in Likelihood.__call__(self, input, *args, **kwargs)
    356 # Marginal
    357 elif any(
    358     [
    359         isinstance(input, MultivariateNormal),
   (...)
    365     ]
    366 ):
--> 367     return self.marginal(input, *args, **kwargs)  # pyre-ignore[6]
    368 # Error
    369 else:
    370     raise RuntimeError(
    371         \"Likelihoods expects a MultivariateNormal or Normal input to make marginal predictions, or a \"
    372         \"torch.Tensor for conditional predictions. Got a {}\".format(input.__class__.__name__)
    373     )