MCPcopy
hub / github.com/dmlc/dgl / __call__

Method __call__

python/dgl/sampling/utils.py:65–105  ·  view source on GitHub ↗
(self, frontier, weights=None)

Source from the content-addressed store, hash-verified

63 return recursive_apply_pair(self._filter, parent_eids, func)
64
65 def __call__(self, frontier, weights=None):
66 parent_eids = frontier.edata[EID]
67 located_eids = self._find_indices(parent_eids)
68
69 if not isinstance(located_eids, Mapping):
70 # (BarclayII) If frontier already has a EID field and located_eids is empty,
71 # the returned graph will keep EID intact. Otherwise, EID will change
72 # to the mapping from the new graph to the old frontier.
73 # So we need to test if located_eids is empty, and do the remapping ourselves.
74 if len(located_eids) > 0:
75 frontier = transforms.remove_edges(
76 frontier, located_eids, store_ids=True
77 )
78 if (
79 weights is not None
80 and weights[0].shape[0] == frontier.num_edges()
81 ):
82 weights[0] = F.gather_row(weights[0], frontier.edata[EID])
83 frontier.edata[EID] = F.gather_row(
84 parent_eids, frontier.edata[EID]
85 )
86 else:
87 # (BarclayII) remove_edges only accepts removing one type of edges,
88 # so I need to keep track of the edge IDs left one by one.
89 new_eids = parent_eids.copy()
90 for i, (k, v) in enumerate(located_eids.items()):
91 if len(v) > 0:
92 frontier = transforms.remove_edges(
93 frontier, v, etype=k, store_ids=True
94 )
95 new_eids[k] = F.gather_row(
96 parent_eids[k], frontier.edges[k].data[EID]
97 )
98 if weights is not None and weights[i].shape[
99 0
100 ] == frontier.num_edges(k):
101 weights[i] = F.gather_row(
102 weights[i], frontier.edges[k].data[EID]
103 )
104 frontier.edata[EID] = new_eids
105 return frontier if weights is None else (frontier, weights)

Callers

nothing calls this directly

Calls 4

_find_indicesMethod · 0.95
remove_edgesMethod · 0.80
num_edgesMethod · 0.45
itemsMethod · 0.45

Tested by

no test coverage detected