Replace length_safe_zip with zip(..., strict=True)

Author: SebastianAment

gpytorch/kernels/lcm_kernel.py

@@ -7,7 +7,6 @@
 from torch.nn import ModuleList
 
 from ..priors import Prior
-from ..utils.generic import length_safe_zip
 from .kernel import Kernel
 from .multitask_kernel import MultitaskKernel
 
@@ -50,7 +49,7 @@ def __init__(
         self.covar_module_list = ModuleList(
             [
                 MultitaskKernel(base_kernel, num_tasks=num_tasks, rank=r, task_covar_prior=task_covar_prior)
-                for base_kernel, r in length_safe_zip(base_kernels, rank)
+                for base_kernel, r in zip(base_kernels, rank, strict=True)
             ]
         )
 

gpytorch/likelihoods/likelihood_list.py

@@ -3,7 +3,6 @@
 from torch.nn import ModuleList
 
 from gpytorch.likelihoods import Likelihood
-from gpytorch.utils.generic import length_safe_zip
 
 
 def _get_tuple_args_(*args):
@@ -22,7 +21,7 @@ def __init__(self, *likelihoods):
     def expected_log_prob(self, *args, **kwargs):
         return [
             likelihood.expected_log_prob(*args_, **kwargs)
-            for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+            for likelihood, args_ in zip(self.likelihoods, _get_tuple_args_(*args), strict=True)
         ]
 
     def forward(self, *args, **kwargs):
@@ -31,18 +30,18 @@ def forward(self, *args, **kwargs):
             # 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)
+                for likelihood, args_, noise_ in zip(self.likelihoods, _get_tuple_args_(*args), noise, strict=True)
             ]
         else:
             return [
                 likelihood.forward(*args_, **kwargs)
-                for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+                for likelihood, args_ in zip(self.likelihoods, _get_tuple_args_(*args), strict=True)
             ]
 
     def pyro_sample_output(self, *args, **kwargs):
         return [
             likelihood.pyro_sample_output(*args_, **kwargs)
-            for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+            for likelihood, args_ in zip(self.likelihoods, _get_tuple_args_(*args), strict=True)
         ]
 
     def __call__(self, *args, **kwargs):
@@ -51,10 +50,10 @@ def __call__(self, *args, **kwargs):
             # if noise kwarg is passed, assume it's an iterable of noise tensors
             return [
                 likelihood(*args_, {**kwargs, "noise": noise_})
-                for likelihood, args_, noise_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args), noise)
+                for likelihood, args_, noise_ in zip(self.likelihoods, _get_tuple_args_(*args), noise, strict=True)
             ]
         else:
             return [
                 likelihood(*args_, **kwargs)
-                for likelihood, args_ in length_safe_zip(self.likelihoods, _get_tuple_args_(*args))
+                for likelihood, args_ in zip(self.likelihoods, _get_tuple_args_(*args), strict=True)
             ]

gpytorch/mlls/sum_marginal_log_likelihood.py

@@ -3,7 +3,6 @@
 from torch.nn import ModuleList
 
 from gpytorch.mlls import ExactMarginalLogLikelihood, MarginalLogLikelihood
-from gpytorch.utils.generic import length_safe_zip
 
 
 class SumMarginalLogLikelihood(MarginalLogLikelihood):
@@ -31,10 +30,10 @@ def forward(self, outputs, targets, *params):
                 (e.g. parameters in case of heteroskedastic likelihoods)
         """
         if len(params) == 0:
-            sum_mll = sum(mll(output, target) for mll, output, target in length_safe_zip(self.mlls, outputs, targets))
+            sum_mll = sum(mll(output, target) for mll, output, target in zip(self.mlls, outputs, targets, strict=True))
         else:
             sum_mll = sum(
                 mll(output, target, *iparams)
-                for mll, output, target, iparams in length_safe_zip(self.mlls, outputs, targets, params)
+                for mll, output, target, iparams in zip(self.mlls, outputs, targets, params, strict=True)
             )
         return sum_mll.div_(len(self.mlls))

gpytorch/models/exact_gp.py

@@ -16,7 +16,6 @@
 from .. import settings
 from ..distributions import MultitaskMultivariateNormal, MultivariateNormal
 from ..likelihoods import _GaussianLikelihoodBase
-from ..utils.generic import length_safe_zip
 from ..utils.warnings import GPInputWarning
 from .exact_prediction_strategies import prediction_strategy
 from .gp import GP
@@ -129,7 +128,7 @@ def set_train_data(
                 inputs = (inputs,)
             inputs = tuple(input_.unsqueeze(-1) if input_.ndimension() == 1 else input_ for input_ in inputs)
             if strict:
-                for input_, t_input in length_safe_zip(inputs, self.train_inputs or (None,)):
+                for input_, t_input in zip(inputs, self.train_inputs or (None,), strict=True):
                     for attr in {"shape", "dtype", "device"}:
                         expected_attr = getattr(t_input, attr, None)
                         found_attr = getattr(input_, attr, None)
@@ -222,7 +221,7 @@ def get_fantasy_model(self, inputs, targets, **kwargs):
                 [train_input, input.expand(input_batch_shape + input.shape[-2:])],
                 dim=-2,
             )
-            for train_input, input in length_safe_zip(train_inputs, inputs)
+            for train_input, input in zip(train_inputs, inputs, strict=True)
         ]
         full_targets = torch.cat(
             [train_targets, targets.expand(target_batch_shape + targets.shape[data_dim_start:])], dim=data_dim_start
@@ -277,7 +276,7 @@ def __call__(self, *args, **kwargs):
                 )
             if settings.debug.on():
                 if not all(
-                    torch.equal(train_input, input) for train_input, input in length_safe_zip(train_inputs, inputs)
+                    torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs, strict=True)
                 ):
                     raise RuntimeError("You must train on the training inputs!")
             res = super().__call__(*inputs, **kwargs)
@@ -295,7 +294,7 @@ def __call__(self, *args, **kwargs):
         # Posterior mode
         else:
             if settings.debug.on():
-                if all(torch.equal(train_input, input) for train_input, input in length_safe_zip(train_inputs, inputs)):
+                if all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs, strict=True)):
                     warnings.warn(
                         "The input matches the stored training data. Did you forget to call model.train()?",
                         GPInputWarning,
@@ -381,7 +380,7 @@ def _get_test_prior_mean_and_covariances(
         # Concatenate the input to the training input
         full_inputs = []
         batch_shape = train_inputs[0].shape[:-2]
-        for train_input, input in length_safe_zip(train_inputs, test_inputs):
+        for train_input, input in zip(train_inputs, test_inputs, strict=True):
             # Make sure the batch shapes agree for training/test data
             if batch_shape != train_input.shape[:-2]:
                 batch_shape = torch.broadcast_shapes(batch_shape, train_input.shape[:-2])

gpytorch/models/model_list.py

@@ -7,7 +7,6 @@
 
 from gpytorch.likelihoods import LikelihoodList
 from gpytorch.models import GP
-from gpytorch.utils.generic import length_safe_zip
 
 
 class AbstractModelList(GP, ABC):
@@ -39,7 +38,7 @@ def likelihood_i(self, i, *args, **kwargs):
 
     def forward(self, *args, **kwargs):
         return [
-            model.forward(*args_, **kwargs) for model, args_ in length_safe_zip(self.models, _get_tensor_args(*args))
+            model.forward(*args_, **kwargs) for model, args_ in zip(self.models, _get_tensor_args(*args), strict=True)
         ]
 
     def get_fantasy_model(self, inputs, targets, **kwargs):
@@ -66,18 +65,19 @@ def get_fantasy_model(self, inputs, targets, **kwargs):
 
         fantasy_models = [
             model.get_fantasy_model(*inputs_, *targets_, **kwargs_)
-            for model, inputs_, targets_, kwargs_ in length_safe_zip(
+for model, inputs_, targets_, kwargs_ in zip(
                 self.models,
                 _get_tensor_args(*inputs),
                 _get_tensor_args(*targets),
                 kwargs,
+                strict=True,
             )
         ]
         return self.__class__(*fantasy_models)
 
     def __call__(self, *args, **kwargs):
         return [
-            model.__call__(*args_, **kwargs) for model, args_ in length_safe_zip(self.models, _get_tensor_args(*args))
+            model.__call__(*args_, **kwargs) for model, args_ in zip(self.models, _get_tensor_args(*args), strict=True)
         ]
 
     @property

gpytorch/utils/generic.py

@@ -2,18 +2,6 @@
 
 from __future__ import annotations
 
-
-def length_safe_zip(*args):
-    """Python's `zip(...)` with checks to ensure the arguments have
-    the same number of elements.
-
-    NOTE: This converts all args that do not define "__len__" to a list.
-    """
-    args = [a if hasattr(a, "__len__") else list(a) for a in args]
-    if len({len(a) for a in args}) > 1:
-        raise ValueError(
-            "Expected the lengths of all arguments to be equal. Got lengths "
-            f"{[len(a) for a in args]} for args {args}. Did you pass in "
-            "fewer inputs than expected?"
-        )
-    return zip(*args)
+# This module previously contained the `length_safe_zip` function.
+# That function has been removed in favor of using Python's built-in
+# `zip(..., strict=True)` which provides the same length-checking behavior.