Invoke user-defined node function on the given nodes. Parameters ---------- graph : DGLGraph The input graph. nid : Tensor The IDs of the nodes to invoke UDF on. ntype : str Node type. func : callable The user-defined function. ndata : dic
(graph, nid, ntype, func, *, ndata=None, orig_nid=None)
| 14 | |
| 15 | |
| 16 | def invoke_node_udf(graph, nid, ntype, func, *, ndata=None, orig_nid=None): |
| 17 | """Invoke user-defined node function on the given nodes. |
| 18 | |
| 19 | Parameters |
| 20 | ---------- |
| 21 | graph : DGLGraph |
| 22 | The input graph. |
| 23 | nid : Tensor |
| 24 | The IDs of the nodes to invoke UDF on. |
| 25 | ntype : str |
| 26 | Node type. |
| 27 | func : callable |
| 28 | The user-defined function. |
| 29 | ndata : dict[str, Tensor], optional |
| 30 | If provided, apply the UDF on this ndata instead of the ndata of the graph. |
| 31 | orig_nid : Tensor, optional |
| 32 | Original node IDs. Useful if the input graph is an extracted subgraph. |
| 33 | |
| 34 | Returns |
| 35 | ------- |
| 36 | dict[str, Tensor] |
| 37 | Results from running the UDF. |
| 38 | """ |
| 39 | ntid = graph.get_ntype_id(ntype) |
| 40 | if ndata is None: |
| 41 | if is_all(nid): |
| 42 | ndata = graph._node_frames[ntid] |
| 43 | nid = graph.nodes(ntype=ntype) |
| 44 | else: |
| 45 | ndata = graph._node_frames[ntid].subframe(nid) |
| 46 | nbatch = NodeBatch( |
| 47 | graph, nid if orig_nid is None else orig_nid, ntype, ndata |
| 48 | ) |
| 49 | return func(nbatch) |
| 50 | |
| 51 | |
| 52 | def invoke_edge_udf(graph, eid, etype, func, *, orig_eid=None): |
no test coverage detected