Function extract_node_subframes

python/dgl/utils/internal.py:842–880

Extract node features of the given nodes from :attr:`graph` and return them in frames on the given device. Note that this function does not perform actual tensor memory copy but using `Frame.subframe` to get the features. If :attr:`nodes` is None, it performs a shallow copy of the o

(graph, nodes_or_device, store_ids=True)

def extract_node_subframes(graph, nodes_or_device, store_ids=True):
    """Extract node features of the given nodes from :attr:`graph`
    and return them in frames on the given device.

    Note that this function does not perform actual tensor memory copy but using `Frame.subframe`
    to get the features. If :attr:`nodes` is None, it performs a shallow copy of the
    original node frames that only copies the dictionary structure but not the tensor
    contents.

    Parameters
    ----------
    graph : DGLGraph
        The graph to extract features from.
    nodes : list[Tensor] or device or None
        Node IDs or device.
        If a list, the list length must be equal to the number of node types
        in the graph.
        If None, the whole frame is shallow-copied.
    store_ids : bool
        If True, the returned frames will store :attr:`nodes` in the ``dgl.NID`` field
        unless it is None.

    Returns
    -------
    list[Frame]
        Extracted node frames.
    """
    if nodes_or_device is None:
        node_frames = [nf.clone() for nf in graph._node_frames]
    elif is_listlike(nodes_or_device):
        node_frames = []
        for i, ind_nodes in enumerate(nodes_or_device):
            subf = graph._node_frames[i].subframe(ind_nodes)
            if store_ids:
                subf[NID] = ind_nodes
            node_frames.append(subf)
    else:  # device object
        node_frames = [nf.to(nodes_or_device) for nf in graph._node_frames]
    return node_frames


def extract_node_subframes_for_block(graph, srcnodes, dstnodes):