move SharedBitGen to random namespace

google · Jan 6, 2024 · adf9e81 · adf9e81
1 parent 49afc4c
commit adf9e81
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 20 deletions.
diff --git a/src/trainer_interface.cc b/src/trainer_interface.cc
@@ -303,8 +303,8 @@ bool TrainerInterface::IsValidSentencePiece(
 }
 
 template <typename T>
-void AddDPNoise(const TrainerSpec &trainer_spec, absl::SharedBitGen &generator,
-                T *to_update) {
+void AddDPNoise(const TrainerSpec &trainer_spec,
+                random::SharedBitGen &generator, T *to_update) {
   if (trainer_spec.differential_privacy_noise_level() > 0) {
     float random_num = absl::Gaussian<float>(
         generator, 0, trainer_spec.differential_privacy_noise_level());
@@ -480,7 +480,7 @@ util::Status TrainerInterface::LoadSentences() {
       for (int n = 0; n < num_workers; ++n) {
         pool->Schedule([&, n]() {
           // One per thread generator.
-          absl::SharedBitGen generator;
+          random::SharedBitGen generator;
           for (size_t i = n; i < sentences_.size(); i += num_workers) {
             AddDPNoise<int64>(trainer_spec_, generator,
                               &(sentences_[i].second));

diff --git a/src/util.h b/src/util.h
@@ -288,6 +288,11 @@ namespace random {
 
 std::mt19937 *GetRandomGenerator();
 
+class SharedBitGen {
+ public:
+  std::mt19937 *engine() { return GetRandomGenerator(); }
+};
+
 template <typename T>
 class ReservoirSampler {
  public:

diff --git a/third_party/absl/random/distributions.h b/third_party/absl/random/distributions.h
@@ -21,8 +21,8 @@
 
 namespace absl {
 
-template <typename T>
-T Gaussian(SharedBitGen &generator, T mean, T stddev) {
+template <typename T, typename G>
+T Gaussian(G &generator, T mean, T stddev) {
   std::normal_distribution<> dist(mean, stddev);
   return dist(*generator.engine());
 }

diff --git a/third_party/absl/random/random.h b/third_party/absl/random/random.h
@@ -15,19 +15,4 @@
 #ifndef ABSL_CONTAINER_RANDOM_H_
 #define ABSL_CONTAINER_RANDOM_H_
 
-#include <random>
-
-#include "../../../src/util.h"
-
-using sentencepiece::random::GetRandomGenerator;
-
-namespace absl {
-
-class SharedBitGen {
- public:
-  std::mt19937 *engine() { return GetRandomGenerator(); }
-};
-
-}  // namespace absl
-
 #endif  // ABSL_CONTAINER_RANDOM_H_