diff --git a/aepsych/models/base.py b/aepsych/models/base.py index 7fbedf9b1..d5d564253 100644 --- a/aepsych/models/base.py +++ b/aepsych/models/base.py @@ -116,7 +116,7 @@ class AEPsychMixin(GPyTorchModel): extremum_solver = "Nelder-Mead" outcome_types: List[str] = [] - train_inputs: Optional[Tuple[torch.Tensor]] + train_inputs: Optional[Tuple[torch.Tensor, ...]] train_targets: Optional[torch.Tensor] @property @@ -393,7 +393,7 @@ def p_below_threshold( class AEPsychModelDeviceMixin(AEPsychMixin): - _train_inputs: Optional[Tuple[torch.Tensor]] + _train_inputs: Optional[Tuple[torch.Tensor, ...]] _train_targets: Optional[torch.Tensor] def set_train_data(self, inputs=None, targets=None, strict=False): @@ -423,13 +423,17 @@ def device(self) -> torch.device: return torch.device("cpu") @property - def train_inputs(self) -> Optional[Tuple[torch.Tensor]]: + def train_inputs(self) -> Optional[Tuple[torch.Tensor, ...]]: if self._train_inputs is None: return None # makes sure the tensors are on the right device, move in place + _train_inputs = [] for input in self._train_inputs: - input.to(self.device) + _train_inputs.append(input.to(self.device)) + + _tuple_inputs: Tuple[torch.Tensor, ...] = tuple(_train_inputs) + self._train_inputs = _tuple_inputs return self._train_inputs diff --git a/tests_gpu/models/test_derivative_gp.py b/tests_gpu/models/test_derivative_gp.py new file mode 100644 index 000000000..200ef62eb --- /dev/null +++ b/tests_gpu/models/test_derivative_gp.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from aepsych import Config, SequentialStrategy +from aepsych.models.derivative_gp import MixedDerivativeVariationalGP +from botorch.fit import fit_gpytorch_mll +from botorch.utils.testing import BotorchTestCase +from gpytorch.likelihoods import BernoulliLikelihood +from gpytorch.mlls.variational_elbo import VariationalELBO + + +class TestDerivativeGP(BotorchTestCase): + def test_MixedDerivativeVariationalGP_gpu(self): + train_x = torch.cat( + (torch.tensor([1.0, 2.0, 3.0, 4.0]).unsqueeze(1), torch.zeros(4, 1)), dim=1 + ) + train_y = torch.tensor([1.0, 2.0, 3.0, 4.0]) + m = MixedDerivativeVariationalGP( + train_x=train_x, + train_y=train_y, + inducing_points=train_x, + fixed_prior_mean=0.5, + ).cuda() + + self.assertEqual(m.mean_module.constant.item(), 0.5) + self.assertEqual( + m.covar_module.base_kernel.raw_lengthscale.shape, torch.Size([1, 1]) + ) + mll = VariationalELBO( + likelihood=BernoulliLikelihood(), model=m, num_data=train_y.numel() + ).cuda() + mll = fit_gpytorch_mll(mll) + test_x = torch.tensor([[1.0, 0], [3.0, 1.0]]).cuda() + m(test_x)