[Bug] `protein.mat` dataset used in "Exact GP Regression with Multiple GPUs" example does not exist.

[rcasero]

🐛 Bug

protein.mat dataset used in "Exact GP Regression with Multiple GPUs" example does not exist.

To reproduce

** Code snippet to reproduce **

The "Exact GP Regression with Multiple GPUs" example uses this bit of code to download the example data from a Google drive

import os
import urllib.request
from scipy.io import loadmat
dataset = 'protein'
if not os.path.isfile(f'../{dataset}.mat'):
    print(f'Downloading \'{dataset}\' UCI dataset...')
    urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1nRb8e7qooozXkNghC5eQS0JeywSXGX2S',
                               f'../{dataset}.mat')

However, just browsing that URL (https://drive.google.com/uc?export=download&id=1nRb8e7qooozXkNghC5eQS0JeywSXGX2S) shows a 404 error. It seems that the document was deleted.

** Stack trace/error message **

N/A

Expected Behavior

The dataset should exist and be downloadable.

System information

Please complete the following information:

• 1.15.1
• 2.0.1+cu117
• Linux 6.17.0-8-generic ImportError: No module named 'gpytorch.random_variables' #8-Ubuntu

Additional context

Add any other context about the problem here.

[kayween]

There have been issues with some other datasets as well in the past; see #1636. We definitely need to update the notebook / google drive.

To unblock you, you can load the dataset from other sources, e.g., https://github.com/treforevans/uci_datasets

[rcasero]

@kayween I forked to solve the problem with the non-existing dataset, but after fixing that I have found the same error that I reported in #2700. MultiGPU doesn't work for me in this example either.

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:1)

My forked version of the "Exact GP Regression with Multiple GPUs" example is here:

https://github.com/rcasero/gpytorch/blob/2699-bug-protein-dataset-url-does-not-exist/examples/02_Scalable_Exact_GPs/Simple_MultiGPU_GP_Regression.ipynb

I have 4 GPUs. When I try to run it and get to this bit of code:

model, likelihood = train(train_x, train_y,
                          n_devices=n_devices, output_device=output_device,
                          preconditioner_size=100,
                          n_training_iter=20)

the full error is

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In [6], [line 1](vscode-notebook-cell:?execution_count=6&line=1)
----> [1](vscode-notebook-cell:?execution_count=6&line=1) model, likelihood = train(train_x, train_y,
      2                           n_devices=n_devices, output_device=output_device,
      3                           preconditioner_size=100,
      4                           n_training_iter=20)

Cell In [5], [line 42](vscode-notebook-cell:?execution_count=5&line=42)
     39     loss = -mll(output, train_y)
     40     return loss
---> [42](vscode-notebook-cell:?execution_count=5&line=42) loss = closure()
     43 loss.backward()
     45 for i in range(n_training_iter):

Cell In [5], [line 39](vscode-notebook-cell:?execution_count=5&line=39)
     37 optimizer.zero_grad()
     38 output = model(train_x)
---> [39](vscode-notebook-cell:?execution_count=5&line=39) loss = -mll(output, train_y)
     40 return loss

File ~/Software/gpytorch/gpytorch/module.py:82, in Module.__call__(self, *inputs, **kwargs)
     81 def __call__(self, *inputs, **kwargs) -> Union[Tensor, Distribution, LinearOperator]:
---> [82](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/gpytorch/module.py:82)     outputs = self.forward(*inputs, **kwargs)
     83     if isinstance(outputs, list):
     84         return [_validate_module_outputs(output) for output in outputs]

File ~/Software/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py:82, in ExactMarginalLogLikelihood.forward(self, function_dist, target, *params, **kwargs)
     79     raise ValueError("NaN observation policy 'fill' is not supported by ExactMarginalLogLikelihood!")
     81 # Get the log prob of the marginal distribution
---> [82](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py:82) res = output.log_prob(target)
     83 res = self._add_other_terms(res, params)
     85 # Scale by the amount of data we have

File ~/Software/gpytorch/gpytorch/distributions/multivariate_normal.py:250, in MultivariateNormal.log_prob(self, value)
    248 # Get log determininant and first part of quadratic form
    249 covar = covar.evaluate_kernel()
--> [250](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/gpytorch/distributions/multivariate_normal.py:250) inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    252 res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])
    253 return res

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:1756, in LinearOperator.inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
   1753 if inv_quad_rhs is not None:
   1754     args = [inv_quad_rhs] + list(args)
-> [1756](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:1756) preconditioner, precond_lt, logdet_p = self._preconditioner()
   1757 if precond_lt is None:
   1758     from linear_operator.operators.identity_linear_operator import IdentityLinearOperator

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/added_diag_linear_operator.py:126, in AddedDiagLinearOperator._preconditioner(self)
    124 if self._q_cache is None:
    125     max_iter = settings.max_preconditioner_size.value()
--> [126](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/added_diag_linear_operator.py:126)     self._piv_chol_self = self._linear_op.pivoted_cholesky(rank=max_iter)
    127     if torch.any(torch.isnan(self._piv_chol_self)).item():
    128         warnings.warn(
    129             "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.",
    130             NumericalWarning,
    131         )

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:1973, in LinearOperator.pivoted_cholesky(self, rank, error_tol, return_pivots)
   1952 r"""
   1953 Performs a partial pivoted Cholesky factorization of the (positive definite) LinearOperator.
   1954 :math:`\mathbf L \mathbf L^\top = \mathbf K`.
   (...)
   1970     https://www.sciencedirect.com/science/article/pii/S0168927411001814
   1971 """
   1972 func = PivotedCholesky.apply
-> [1973](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:1973) res, pivots = func(self.representation_tree(), rank, error_tol, *self.representation())
   1975 if return_pivots:
   1976     return res, pivots

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/torch/autograd/function.py:583, in Function.apply(cls, *args, **kwargs)
    580 if not torch._C._are_functorch_transforms_active():
    581     # See NOTE: [functorch vjp and autograd interaction]
    582     args = _functorch.utils.unwrap_dead_wrappers(args)
--> [583](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/torch/autograd/function.py:583)     return super().apply(*args, **kwargs)  # type: ignore[misc]
    585 if not is_setup_ctx_defined:
    586     raise RuntimeError(
    587         "In order to use an autograd.Function with functorch transforms "
    588         "(vmap, grad, jvp, jacrev, ...), it must override the setup_context "
    589         "staticmethod. For more details, please see "
    590         "https://pytorch.org/docs/main/notes/extending.func.html"
    591     )

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/functions/_pivoted_cholesky.py:78, in PivotedCholesky.forward(ctx, representation_tree, max_iter, error_tol, *matrix_args)
     75 # Populater L[... m:, m] with L[..., m:, m] * L[..., m, m].sqrt()
     76 if m + 1 < matrix_shape[-1]:
     77     # Get next row of the permuted matrix
---> [78](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/functions/_pivoted_cholesky.py:78)     row = apply_permutation(matrix, pi_m.unsqueeze(-1), right_permutation=None).squeeze(-2)
     79     pi_i = permutation[..., m + 1 :].contiguous()
     81     L_m_new = row.gather(-1, pi_i)

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/utils/permutation.py:80, in apply_permutation(matrix, left_permutation, right_permutation)
     76     right_permutation = torch.arange(matrix.size(-1), device=matrix.device)
     78 # Apply permutations
     79 return to_dense(
---> [80](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/utils/permutation.py:80)     matrix.__getitem__(
     81         (
     82             *batch_idx,
     83             left_permutation.unsqueeze(-1),
     84             right_permutation.unsqueeze(-2),
     85         )
     86     )
     87 )

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:2855, in LinearOperator.__getitem__(self, index)
   2849 # Convert all indices into tensor indices
   2850 (
   2851     *new_batch_indices,
   2852     new_row_index,
   2853     new_col_index,
   2854 ) = _convert_indices_to_tensors(self, flattened_orig_indices)
-> [2855](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/_linear_operator.py:2855) res = self._get_indices(new_row_index, new_col_index, *new_batch_indices)
   2856 # Now un-flatten tensor indices
   2857 if len(tensor_index_shape) > 1:  # Do we need to unflatten?

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/cat_linear_operator.py:214, in CatLinearOperator._get_indices(self, row_index, col_index, *batch_indices)
    210 for linear_op_idx, sub_index in zip(linear_op_indices, sub_indices):
    211     sub_index[self.cat_dim] = sub_index[self.cat_dim] - self.cat_dim_cum_sizes[linear_op_idx]
    213 res_list = [
--> [214](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/cat_linear_operator.py:214)     linear_op._get_indices(sub_index[-2], sub_index[-1], *sub_index[:-2])
    215     for linear_op, sub_index in zip(linear_ops, sub_indices)
    216 ]
    217 if len(res_list) == 1:
    218     return res_list[0].view(target_shape).to(self.device)

File ~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/dense_linear_operator.py:50, in DenseLinearOperator._get_indices(self, row_index, col_index, *batch_indices)
     48 def _get_indices(self, row_index: IndexType, col_index: IndexType, *batch_indices: IndexType) -> torch.Tensor:
     49     # Perform the __getitem__
---> [50](https://vscode-remote+ssh-002dremote-002bbeast2.vscode-resource.vscode-cdn.net/home/ramon/Software/gpytorch/examples/02_Scalable_Exact_GPs/~/Software/gpytorch/.venv/lib/python3.13/site-packages/linear_operator/operators/dense_linear_operator.py:50)     res = self.tensor[(*batch_indices, row_index, col_index)]
     51     return res

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:1)

[kayween]

@rcasero you could try the keops kernel (a memery efficient implementation of dense kernel matrices), while we are working on fixing those bugs.

https://docs.gpytorch.ai/en/stable/examples/02_Scalable_Exact_GPs/KeOps_GP_Regression.html

It seems like you have about 100_000 data points. KeOps kernels should be able to handle that on a single GPU.

[rcasero]

@kayween Wonderful advice, many thanks for the KeOps pointer. I've tried it on reduced data, and it works. I had some trouble initially, because if you select a KeOps kernel but have not installed the dependency with gpytorch[keops], the non-KeOps kernel is silently used. I'll report if I get it to work on full data (I have multiple labels and also ~100 features).

[rcasero]

@kayween I have created a separate issue #2702 just for the multi-GPU, so that this issue can be closed after fixing the problem with the missing protein dataset. Issue #2700 can be referred to the new #2702 issue, as once the multi-GPU support is fixed, then combining it with multiple label classification will probably work too.

[kayween]

@jacobrgardner @gpleiss The google drive hosting the UCI dataset has been corrupted for a while. It would be great to fix this so that users can run example notebooks.

We could figure out what's going on with the google drive and and update the link.

The other option is to rely on external packages to load the data. I have been using this repo for GP benchmarking and it works well: https://github.com/treforevans/uci_datasets I am leaning towards this and and I am wondering if you guys have strong opinions on this.

The pro of using external packages for data loading is that it releases us from the maintenance burden. The con is that it introduces an extra dependency, and the above repo in particular is not super active.