Source code for jittor_geometric.utils.num_nodes

from jittor import Var




def maybe_num_nodes(edge_index, num_nodes=None):
    if num_nodes is not None:
        return num_nodes
    elif isinstance(edge_index, Var):
        return int(edge_index.max()) + 1
    else:
        return max(edge_index.size(0), edge_index.size(1))