Skip to content

Commit ae22058

Browse files
yaoyaowdDong Wang
and
Dong Wang
authored
make graph sample deterministic when iterate different type (#265)
Co-authored-by: Dong Wang <[email protected]>
1 parent 7f6379a commit ae22058

File tree

1 file changed

+8
-3
lines changed

1 file changed

+8
-3
lines changed

csrc/cpu/neighbor_sample_cpu.cpp

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,13 +192,18 @@ hetero_sample(const vector<node_t> &node_types,
192192
for (const auto &kv : samples_dict)
193193
slice_dict[kv.first] = {0, kv.second.size()};
194194

195+
vector<rel_t> all_rel_types;
196+
for (const auto &kv : num_neighbors_dict) {
197+
all_rel_types.push_back(kv.key());
198+
}
199+
std::sort(all_rel_types.begin(), all_rel_types.end());
200+
195201
for (int64_t ell = 0; ell < num_hops; ell++) {
196-
for (const auto &kv : num_neighbors_dict) {
197-
const auto &rel_type = kv.key();
202+
for (const auto &rel_type : all_rel_types) {
198203
const auto &edge_type = to_edge_type[rel_type];
199204
const auto &src_node_type = get<0>(edge_type);
200205
const auto &dst_node_type = get<2>(edge_type);
201-
const auto num_samples = kv.value()[ell];
206+
const auto num_samples = num_neighbors_dict.at(rel_type)[ell];
202207
const auto &dst_samples = samples_dict.at(dst_node_type);
203208
auto &src_samples = samples_dict.at(src_node_type);
204209
auto &to_local_src_node = to_local_node_dict.at(src_node_type);

0 commit comments

Comments
 (0)