c517e0b8b34a85b8142b4669c152b6e62c02d8e2,gpytorch/lazy/lazy_tensor.py,LazyTensor,diag,#LazyTensor#,567

Before Change

            raise RuntimeError("Diag works on square matrices (or batches)")

        row_col_iter = torch.arange(0, size[-1], dtype=torch.long, device=self.device)
        if self.ndimension() == 3:
            batch_iter = torch.arange(0, size[0], dtype=torch.long, device=self.device)
            batch_iter = batch_iter.unsqueeze(1).repeat(1, size[1]).view(-1)
            row_col_iter = row_col_iter.unsqueeze(1).repeat(size[0], 1).view(-1)
            return self._get_indices(row_col_iter, row_col_iter, batch_iter).view(size[0], size[1])
        else:
            return self._get_indices(row_col_iter, row_col_iter)

    def dim(self):
        
        Alias of :meth:`~gpytorch.lazy.LazyTensor.ndimension`
        

After Change

            n vector. If this LazyTensor represents a batch (e.g., is :math:`b \times n \times n`), this will be a
            :math:`b \times n` matrix of diagonals, one for each matrix in the batch.
        
        if settings.debug.on():
            if not self.is_square:
                raise RuntimeError("Diag works on square matrices (or batches)")

        row_col_iter = torch.arange(0, self.matrix_shape[-1], dtype=torch.long, device=self.device)
        return self[..., row_col_iter, row_col_iter]

    def dim(self):

Instances