def weighted_dimwise_median(A: torch.sparse.FloatTensor, x: torch.Tensor,
**kwargs) -> torch.Tensor:
"""A weighted dimension-wise Median aggregation.
Parameters
----------
A : torch.sparse.FloatTensor
Sparse [n, n] tensor of the weighted/normalized adjacency matrix
x : torch.Tensor
Dense [n, d] tensor containing the node attributes/embeddings
Returns
-------
torch.Tensor
The new embeddings [n, d]
"""
if not A.is_cuda:
return weighted_dimwise_median_cpu(A, x, **kwargs)
assert A.is_sparse
N, D = x.shape
median_idx = custom_cuda_kernels.dimmedian_idx(x, A)
col_idx = torch.arange(D, device=A.device).view(1, -1).expand(N, D)
x_selected = x[median_idx, col_idx]
a_row_sum = torch_scatter.scatter_sum(A._values(), A._indices()[0],
dim=-1).view(-1, 1).expand(N, D)
return a_row_sum * x_selected
def add_eye_sparse_tensor(x: torch.sparse.FloatTensor,
bandwidth: int) -> torch.sparse.FloatTensor:
"""Not used!"""
if not bandwidth:
return x
indices_list = [x._indices()]
values_list = [x._values()]
# Add diagonal
indices_list.append(
torch.arange(0, x.shape[1], dtype=torch.long).repeat(2, 1))
values_list.append(torch.ones(x.shape[1], dtype=torch.float))
# Add off-diagonals
for k in range(1, bandwidth, 1):
# Indices
indices_upper = torch.stack([
torch.arange(x.shape[1] - k, dtype=torch.long),
k + torch.arange(x.shape[1] - k, dtype=torch.long),
])
indices_lower = torch.stack([
k + torch.arange(x.shape[1] - k, dtype=torch.long),
torch.arange(x.shape[1] - k, dtype=torch.long),
])
indices_list.extend([indices_upper, indices_lower])
# Values
values_upper = torch.ones(x.shape[1] - k, dtype=torch.float)
values_lower = torch.ones(x.shape[1] - k, dtype=torch.float)
values_list.extend([values_upper, values_lower])
out = torch.sparse_coo_tensor(torch.cat(indices_list),
torch.cat(values_list),
size=x.size)
return out
def _reduce(x: torch.sparse.FloatTensor):
# dispatch table cannot distinguish between torch.sparse.FloatTensor and torch.Tensor
if isinstance(x, torch.sparse.FloatTensor) or isinstance(
x, torch.sparse.LongTensor):
int_type = _get_int_type(torch.max(x._indices()).item())
return _sparse_tensor_constructor, (x._indices().to(int_type),
x._values(), x.size())
else:
return torch.Tensor.__reduce_ex__(
x, pickle.HIGHEST_PROTOCOL) # use your own protocol
def _sparse_masked_select_abs(self,
sparse_tensor: torch.sparse.FloatTensor,
thr):
indices = sparse_tensor._indices()
values = sparse_tensor._values()
prune_mask = torch.abs(values) >= thr
return torch.sparse_coo_tensor(
indices=indices.masked_select(prune_mask).reshape(2, -1),
values=values.masked_select(prune_mask),
size=[self.n_outputs, self.n_inputs]).coalesce()
def _sparse_top_k(A: torch.sparse.FloatTensor,
k: int,
return_sparse: bool = True):
n = A.shape[0]
if A.is_cuda:
topk_values, topk_idx = custom_cuda_kernels.topk(
A._indices(), A._values(), n, k)
if not return_sparse:
return topk_values, topk_idx.long()
mask = topk_idx != -1
row_idx = torch.arange(n, device=A.device).view(-1, 1).expand(n, k)
return torch.sparse.FloatTensor(
torch.stack((row_idx[mask], topk_idx[mask].long())),
topk_values[mask])
n_edges_per_row = torch_scatter.scatter_sum(torch.ones_like(A._values()),
A._indices()[0],
dim=0)
k_per_row = torch.clamp(n_edges_per_row, max=k).long()
new_idx, value_idx, unroll_idx = _select_k_idx_cpu(
A.indices()[0].cpu().numpy(),
A.indices()[1].cpu().numpy(),
A._values().cpu().numpy(),
k_per_row.cpu().numpy(),
n,
method='top')
new_idx = torch.from_numpy(np.hstack(new_idx)).to(A.device)
value_idx = torch.from_numpy(np.hstack(value_idx)).to(A.device)
if return_sparse:
return torch.sparse.FloatTensor(new_idx, A._values()[value_idx])
else:
unroll_idx = np.hstack(unroll_idx)
values = torch.zeros((n, k), device=A.device)
indices = -torch.ones((n, k), device=A.device, dtype=torch.long)
values[new_idx[0], unroll_idx] = A._values()[value_idx]
indices[new_idx[0], unroll_idx] = new_idx[1]
return values, indices
def remove_eye_sparse_tensor(x: torch.sparse.FloatTensor,
bandwidth: int) -> torch.sparse.FloatTensor:
"""Set diagonal (and offdiagonal) elements to zero.
Args:
x: Input array.
bandwidth: Width of the diagonal 0 band.
"""
if not bandwidth:
return x
indices = x._indices()
values = x._values()
keep_mask = (indices[0, :] - indices[1, :]).abs() > bandwidth
out = torch.sparse_coo_tensor(indices[:, keep_mask], values[keep_mask])
return out