1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
|
#pragma once
#include "../extensions.h"
#include "parallel_hashmap/phmap.h"
#define CHECK_CPU(x) AT_ASSERTM(x.device().is_cpu(), #x " must be CPU tensor")
#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
#define CHECK_LT(low, high) AT_ASSERTM(low < high, "low must be smaller than high")
#define AT_DISPATCH_HAS_VALUE(optional_value, ...) \
[&] { \
if (optional_value.has_value()) { \
const bool HAS_VALUE = true; \
return __VA_ARGS__(); \
} else { \
const bool HAS_VALUE = false; \
return __VA_ARGS__(); \
} \
}()
template <typename scalar_t>
inline torch::Tensor from_vector(const std::vector<scalar_t> &vec,
bool inplace = false) {
const auto size = (int64_t)vec.size();
const auto out = torch::from_blob((scalar_t *)vec.data(), {size},
c10::CppTypeToScalarType<scalar_t>::value);
return inplace ? out : out.clone();
}
template <typename key_t, typename scalar_t>
inline c10::Dict<key_t, torch::Tensor>
from_vector(const phmap::flat_hash_map<key_t, std::vector<scalar_t>> &vec_dict,
bool inplace = false) {
c10::Dict<key_t, torch::Tensor> out_dict;
for (const auto &kv : vec_dict)
out_dict.insert(kv.first, from_vector<scalar_t>(kv.second, inplace));
return out_dict;
}
inline int64_t uniform_randint(int64_t low, int64_t high) {
CHECK_LT(low, high);
auto options = torch::TensorOptions().dtype(torch::kInt64);
auto ret = torch::randint(low, high, {1}, options);
auto ptr = ret.data_ptr<int64_t>();
return *ptr;
}
inline int64_t uniform_randint(int64_t high) {
return uniform_randint(0, high);
}
inline torch::Tensor
choice(int64_t population, int64_t num_samples, bool replace = false,
std::optional<torch::Tensor> weight = std::nullopt) {
if (population == 0 || num_samples == 0)
return torch::empty({0}, at::kLong);
if (!replace && num_samples >= population)
return torch::arange(population, at::kLong);
if (weight.has_value())
return torch::multinomial(weight.value(), num_samples, replace);
if (replace) {
const auto out = torch::empty({num_samples}, at::kLong);
auto *out_data = out.data_ptr<int64_t>();
for (int64_t i = 0; i < num_samples; i++) {
out_data[i] = uniform_randint(population);
}
return out;
} else {
// Sample without replacement via Robert Floyd algorithm:
// https://www.nowherenearithaca.com/2013/05/
// robert-floyds-tiny-and-beautiful.html
const auto out = torch::empty({num_samples}, at::kLong);
auto *out_data = out.data_ptr<int64_t>();
std::unordered_set<int64_t> samples;
for (int64_t i = population - num_samples; i < population; i++) {
int64_t sample = uniform_randint(i);
if (!samples.insert(sample).second) {
sample = i;
samples.insert(sample);
}
out_data[i - population + num_samples] = sample;
}
return out;
}
}
template <bool replace>
inline void
uniform_choice(const int64_t population, const int64_t num_samples,
const int64_t *idx_data, std::vector<int64_t> *samples,
phmap::flat_hash_map<int64_t, int64_t> *to_local_node) {
if (population == 0 || num_samples == 0)
return;
if (replace) {
for (int64_t i = 0; i < num_samples; i++) {
const int64_t &v = idx_data[uniform_randint(population)];
if (to_local_node->insert({v, samples->size()}).second)
samples->push_back(v);
}
} else if (num_samples >= population) {
for (int64_t i = 0; i < population; i++) {
const int64_t &v = idx_data[i];
if (to_local_node->insert({v, samples->size()}).second)
samples->push_back(v);
}
} else {
std::unordered_set<int64_t> indices;
for (int64_t i = population - num_samples; i < population; i++) {
int64_t j = uniform_randint(i);
if (!indices.insert(j).second) {
j = i;
indices.insert(j);
}
const int64_t &v = idx_data[j];
if (to_local_node->insert({v, samples->size()}).second)
samples->push_back(v);
}
}
}
|
