def remove_diag(src: SparseTensor, k: int = 0) -> SparseTensor:
row, col, value = src.coo()
inv_mask = row != col if k == 0 else row != (col - k)
new_row, new_col = row[inv_mask], col[inv_mask]
if value is not None:
value = value[inv_mask]
rowcount = src.storage._rowcount
colcount = src.storage._colcount
if rowcount is not None or colcount is not None:
mask = ~inv_mask
if rowcount is not None:
rowcount = rowcount.clone()
rowcount[row[mask]] -= 1
if colcount is not None:
colcount = colcount.clone()
colcount[col[mask]] -= 1
storage = SparseStorage(row=new_row,
rowptr=None,
col=new_col,
value=value,
sparse_sizes=src.sparse_sizes(),
rowcount=rowcount,
colptr=None,
colcount=colcount,
csr2csc=None,
csc2csr=None,
is_sorted=True)
return src.from_storage(storage)
def set_diag(src: SparseTensor,
values: Optional[torch.Tensor] = None,
k: int = 0) -> SparseTensor:
src = remove_diag(src, k=k)
row, col, value = src.coo()
mask = torch.ops.torch_sparse.non_diag_mask(row, col, src.size(0),
src.size(1), k)
inv_mask = ~mask
start, num_diag = -k if k < 0 else 0, mask.numel() - row.numel()
diag = torch.arange(start, start + num_diag, device=row.device)
new_row = row.new_empty(mask.size(0))
new_row[mask] = row
new_row[inv_mask] = diag
new_col = col.new_empty(mask.size(0))
new_col[mask] = col
new_col[inv_mask] = diag.add_(k)
new_value: Optional[torch.Tensor] = None
if value is not None:
new_value = value.new_empty((mask.size(0), ) + value.size()[1:])
new_value[mask] = value
if values is not None:
new_value[inv_mask] = values
else:
new_value[inv_mask] = torch.ones((num_diag, ),
dtype=value.dtype,
device=value.device)
rowcount = src.storage._rowcount
if rowcount is not None:
rowcount = rowcount.clone()
rowcount[start:start + num_diag] += 1
colcount = src.storage._colcount
if colcount is not None:
colcount = colcount.clone()
colcount[start + k:start + num_diag + k] += 1
storage = SparseStorage(row=new_row,
rowptr=None,
col=new_col,
value=new_value,
sparse_sizes=src.sparse_sizes(),
rowcount=rowcount,
colptr=None,
colcount=colcount,
csr2csc=None,
csc2csr=None,
is_sorted=True)
return src.from_storage(storage)
def masked_select_nnz(src: SparseTensor,
mask: torch.Tensor,
layout: Optional[str] = None) -> SparseTensor:
assert mask.dim() == 1
if get_layout(layout) == 'csc':
mask = mask[src.storage.csc2csr()]
row, col, value = src.coo()
row, col = row[mask], col[mask]
if value is not None:
value = value[mask]
return SparseTensor(row=row,
rowptr=None,
col=col,
value=value,
sparse_sizes=src.sparse_sizes(),
is_sorted=True)
def index_select_nnz(src: SparseTensor,
idx: torch.Tensor,
layout: Optional[str] = None) -> SparseTensor:
assert idx.dim() == 1
if get_layout(layout) == 'csc':
idx = src.storage.csc2csr()[idx]
row, col, value = src.coo()
row, col = row[idx], col[idx]
if value is not None:
value = value[idx]
return SparseTensor(row=row,
rowptr=None,
col=col,
value=value,
sparse_sizes=src.sparse_sizes(),
is_sorted=True)