blob: 4cf58b8979a3becb223c5e046b61519066b636ab (
plain)
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
|
#include "tensorflow/core/util/guarded_philox_random.h"
#include "tensorflow/core/lib/random/random.h"
namespace tensorflow {
Status GuardedPhiloxRandom::Init(OpKernelConstruction* context) {
// Grab seed Attrs.
int64 seed, seed2;
auto status = context->GetAttr("seed", &seed);
if (!status.ok()) return status;
status = context->GetAttr("seed2", &seed2);
if (!status.ok()) return status;
// Initialize with the given seeds
Init(seed, seed2);
return Status::OK();
}
void GuardedPhiloxRandom::Init(int64 seed, int64 seed2) {
CHECK(!initialized_);
if (seed == 0 && seed2 == 0) {
// If both seeds are unspecified, use completely random seeds.
seed = random::New64();
seed2 = random::New64();
}
mutex_lock lock(mu_);
generator_ = random::PhiloxRandom(seed, seed2);
initialized_ = true;
}
random::PhiloxRandom GuardedPhiloxRandom::ReserveSamples128(int64 samples) {
CHECK(initialized_);
mutex_lock lock(mu_);
auto local = generator_;
generator_.Skip(samples);
return local;
}
} // namespace tensorflow
|
