"examples/sampling/vscode:/vscode.git/clone" did not exist on "2fedcdc2feda37b56fcf3198a4eac79ba8ae9a2d"
Commit 258dbf0f authored by rusty1s's avatar rusty1s
Browse files

from_dense with optional no values

parent 6b013fbe
......@@ -30,15 +30,21 @@ class SparseTensor(object):
return self
@classmethod
def from_dense(self, mat: torch.Tensor):
def from_dense(self, mat: torch.Tensor, has_value: bool = True):
if mat.dim() > 2:
index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero()
else:
index = mat.nonzero()
index = index.t()
row, col = index[0], index[1]
return SparseTensor(row=row, rowptr=None, col=col, value=mat[row, col],
row = index[0]
col = index[1]
value: Optional[torch.Tensor] = None
if has_value:
value = mat[row, col]
return SparseTensor(row=row, rowptr=None, col=col, value=value,
sparse_sizes=mat.size()[:2], is_sorted=True)
@classmethod
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment