Source code for torch_geometric.utils.degree

import torch

from .num_nodes import maybe_num_nodes


[docs]def degree(index, num_nodes=None, dtype=None): r"""Computes the (unweighted) degree of a given one-dimensional index tensor. Args: index (LongTensor): Index tensor. num_nodes (int, optional): The number of nodes, *i.e.* :obj:`max_val + 1` of :attr:`index`. (default: :obj:`None`) dtype (:obj:`torch.dtype`, optional): The desired data type of the returned tensor. :rtype: :class:`Tensor` """ num_nodes = maybe_num_nodes(index, num_nodes) out = torch.zeros((num_nodes), dtype=dtype, device=index.device) return out.scatter_add_(0, index, out.new_ones((index.size(0))))