7570ee1d758c93dd4054009fc12be6cc31b54031,gpytorch/lazy/interpolated_lazy_tensor.py,InterpolatedLazyTensor,_get_indices,#InterpolatedLazyTensor#Any#Any#,261

Before Change

            left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
            right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
        else:
            left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
            right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
        base_var_vals = self.base_lazy_tensor._get_indices(left_interp_indices.view(-1), right_interp_indices.view(-1))
        base_var_vals = base_var_vals.view(left_interp_indices.size())
        res = (interp_values * base_var_vals).sum(-1).sum(-1)

After Change

        // Batch compute values that will be non-zero for row k
        left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
        right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
        batch_indices = [
            batch_index.unsqueeze(1).repeat(1, n_interp ** 2).view(-1)
            for batch_index in batch_indices
        ]
        base_var_vals = self.base_lazy_tensor._get_indices(
            left_interp_indices.view(-1), right_interp_indices.view(-1), *batch_indices
        )
        base_var_vals = base_var_vals.view(left_interp_indices.size())

Instances