f76a4dabb4cd38ee58d01a35c5b511e224d060d2,gpytorch/lazy/sum_batch_lazy_tensor.py,SumBatchLazyTensor,_get_indices,#SumBatchLazyTensor#Any#Any#,99

Before Change

        raise RuntimeError("Batch get indices is not implmeneted yet")

    def _get_indices(self, left_indices, right_indices):
        batch_indices = self.tensor_cls(self.batch_size()).long()
        torch.arange(0, self.batch_size(), out=batch_indices.data)
        batch_indices = batch_indices.unsqueeze(1).repeat(1, len(left_indices)).view(-1)
        left_indices = left_indices.unsqueeze(1).repeat(self.batch_size(), 1).view(-1)
        right_indices = right_indices.unsqueeze(1).repeat(self.batch_size(), 1).view(-1)

After Change

        res = res.view(self.num_blocks, -1).sum(0)
        return res

    def _get_indices(self, left_indices, right_indices):
        batch_indices = torch.arange(0, self.base_lazy_tensor.size(0), dtype=torch.long, device=left_indices.device)
        batch_indices = batch_indices.unsqueeze(1).repeat(1, len(left_indices)).view(-1)
        left_indices = left_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
        right_indices = right_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
        return res.view(self.base_lazy_tensor.size(0), -1).sum(0)

Instances