c517e0b8b34a85b8142b4669c152b6e62c02d8e2,gpytorch/lazy/block_diag_lazy_tensor.py,BlockDiagLazyTensor,_size,#BlockDiagLazyTensor#,67

Before Change

    def _size(self):
        base_size = self.base_lazy_tensor.size()
        if self.num_blocks is None:
            return torch.Size((base_size[0] * base_size[1], base_size[0] * base_size[2]))
        else:
            true_batch_size = self.base_lazy_tensor.size(0) // self.num_blocks
            return torch.Size((true_batch_size, self.num_blocks * base_size[1], self.num_blocks * base_size[2]))

    def _get_indices(self, left_indices, right_indices, *batch_indices):
        if self.num_blocks is None:
            if settings.debug.on():
                assert len(batch_indices) == 0

After Change

        shape[-2] *= shape[self.block_dim]
        shape[-1] *= shape[self.block_dim]
        del shape[self.block_dim]
        return torch.Size(shape)

    def diag(self):
        res = self.base_lazy_tensor.diag().contiguous()
        return res.view(*self.batch_shape, self.size(-1))

Instances