range_sampler.h

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_CORE_KERNELS_RANGE_SAMPLER_H_
#define TENSORFLOW_CORE_KERNELS_RANGE_SAMPLER_H_

#include <vector>

#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/lib/random/distribution_sampler.h"
#include "tensorflow/core/lib/random/random_distributions.h"
#include "tensorflow/core/lib/random/weighted_picker.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"

namespace tsl {
class Env;
}  // namespace tsl
namespace tensorflow {
using Env = tsl::Env;

// Abstract subclass for sampling from the set of non-negative integers
// [0, range)
class RangeSampler {
 public:
  explicit RangeSampler(int64_t range) : range_(range) { CHECK_GT(range_, 0); }
  virtual ~RangeSampler();

  // Sample a single value
  virtual int64_t Sample(random::SimplePhilox* rnd) const = 0;

  // The probability that a single call to Sample() returns the given value.
  // Assumes that value is in [0, range).  No range checking is done.
  virtual float Probability(int64_t value) const = 0;

  // Fill "batch" with samples from the distribution.
  // If unique=true, then we re-pick each element until we get a
  // value distinct from all previously picked values in the batch.
  void SampleBatch(random::SimplePhilox* rnd, bool unique,
                   absl::Span<int64_t> batch) const;

  // Fill "batch" with samples from the distribution, and report
  // "expected counts".
  //
  // The "expected count" of a value is an estimate of the expected
  // number of occurrences of the value in the batch returned by a
  // call to this function with the given parameters.  If unique=true,
  // the expected count is an inclusion probability.  For details on
  // this estimation, see the comment to "ExpectedCountHelper" in the
  // .cc file.
  //
  // Expected counts for the elements of the returned "batch" are reported
  // in the aligned array "batch_expected_count".
  //
  // The user can optionally provide "extras", containing values in the range.
  // The expected counts for the extras are reported in the aligned array
  // "extras_expected_count".
  //
  // "batch_expected_count" must have size equal to 0 or to the size of "batch".
  // "extras" and "extras_expected_count" must have equal size.
  void SampleBatchGetExpectedCount(
      random::SimplePhilox* rnd, bool unique, absl::Span<int64_t> batch,
      absl::Span<float> batch_expected_count, absl::Span<const int64_t> extras,
      absl::Span<float> extras_expected_count) const;

  // Same as SampleBatchGetExpectedCount (see above), but with avoided values.
  // We repick to avoid all of the values in "avoided_values".
  // "avoided_values" is only supported with unique=true.  If
  // unique=false, then avoided_values must be empty.
  virtual void SampleBatchGetExpectedCountAvoid(
      random::SimplePhilox* rnd, bool unique, absl::Span<int64_t> batch,
      absl::Span<float> batch_expected_count, absl::Span<const int64_t> extras,
      absl::Span<float> extras_expected_count,
      absl::Span<const int64_t> avoided_values) const;

  // Does this sampler need to be updated with values, e.g. UnigramSampler
  virtual bool NeedsUpdates() const { return false; }

  // Updates the underlying distribution
  virtual void Update(absl::Span<const int64_t> values) {
    LOG(FATAL) << "Update not supported for this sampler type.";
  }

  int64_t range() { return range_; }

 protected:
  const int64_t range_;
};

// An AllSampler only samples batches of size equal to range.
// It returns the entire range.
// It cannot sample single values.
class AllSampler : public RangeSampler {
 public:
  explicit AllSampler(int64_t range);

  ~AllSampler() override {}

  int64_t Sample(random::SimplePhilox* rnd) const override {
    LOG(FATAL) << "Should not be called";
    return 0;
  }

  float Probability(int64_t value) const override {
    LOG(FATAL) << "Should not be called";
    return 0;
  }

  void SampleBatchGetExpectedCountAvoid(
      random::SimplePhilox* rnd, bool unique, absl::Span<int64_t> batch,
      absl::Span<float> batch_expected_count, absl::Span<const int64_t> extras,
      absl::Span<float> extras_expected_count,
      absl::Span<const int64_t> avoided_values) const override;
};

class UniformSampler : public RangeSampler {
 public:
  explicit UniformSampler(int64_t range);

  ~UniformSampler() override {}

  int64_t Sample(random::SimplePhilox* rnd) const override;

  float Probability(int64_t value) const override;

 private:
  const float inv_range_;
};

class LogUniformSampler : public RangeSampler {
 public:
  explicit LogUniformSampler(int64_t range);

  ~LogUniformSampler() override {}

  int64_t Sample(random::SimplePhilox* rnd) const override;

  float Probability(int64_t value) const override;

 private:
  const double log_range_;
};

// Thread-unsafe unigram sampler
class ThreadUnsafeUnigramSampler : public RangeSampler {
 public:
  explicit ThreadUnsafeUnigramSampler(int64_t range);
  ~ThreadUnsafeUnigramSampler() override {}

  int64_t Sample(random::SimplePhilox* rnd) const override;

  float Probability(int64_t value) const override;

  bool NeedsUpdates() const override { return true; }
  void Update(absl::Span<const int64_t> values) override;

 private:
  random::WeightedPicker picker_;
};

// Thread-safe unigram sampler
class UnigramSampler : public RangeSampler {
 public:
  explicit UnigramSampler(int64_t range);
  ~UnigramSampler() override {}

  int64_t Sample(random::SimplePhilox* rnd) const override;

  float Probability(int64_t value) const override;

  // Overriding at a high level results in far fewer lock acquisitions.
  void SampleBatchGetExpectedCountAvoid(
      random::SimplePhilox* rnd, bool unique, absl::Span<int64_t> batch,
      absl::Span<float> batch_expected_count, absl::Span<const int64_t> extras,
      absl::Span<float> extras_expected_count,
      absl::Span<const int64_t> avoided_values) const override;

  bool NeedsUpdates() const override { return true; }
  void Update(absl::Span<const int64_t> values) override;

 private:
  ThreadUnsafeUnigramSampler unsafe_sampler_ TF_GUARDED_BY(mu_);
  mutable mutex mu_;
};

// A unigram sampler that uses a fixed unigram distribution read from a
// file or passed in as an in-memory array instead of building up the
// distribution from data on the fly. There is also an option to skew the
// distribution by applying a distortion power to the weights.
class FixedUnigramSampler : public RangeSampler {
 public:
  FixedUnigramSampler(int64_t range, float distortion, int32_t num_reserved_ids,
                      int32_t num_shards, int32_t shard);
  // The vocab_file is assumed to be a CSV, with the last entry of each row a
  // value representing the counts or probabilities for the corresponding ID.
  Status SetDistributionSampler(Env* env, const string& vocab_file);
  Status SetDistributionSampler(const std::vector<float>& unigrams);
  float Probability(int64_t value) const override;

  int64_t Sample(random::SimplePhilox* rnd) const override;

 private:
  // Underlying distribution sampler.
  std::unique_ptr<random::DistributionSampler> dist_sampler_;
  // Weights for individual samples. The probability of a sample i is defined
  // as weights_.at(i) / total_weight_.
  std::vector<float> weights_;
  // The total weights of all samples.
  float total_weight_;
  // Sharding information of the sampler. The whole vocabulary is sharded
  // into num_shards_ smaller ranges and each sampler is responsible for one
  // such smaller range, identified by the shard number.
  int32 num_shards_;
  int32 shard_;
  float distortion_;
  // Fill the sampler with the appropriate number of reserved IDs.
  void FillReservedIds(int32_t num_reserved_ids);
  // Load IDs to sample from a CSV file. It is assumed that the last item of
  // each row contains a count or probability for the corresponding ID.
  Status LoadFromFile(Env* env, const string& vocab_file, float distortion);
  // Load from an in-memory array.
  void LoadFromUnigrams(const std::vector<float>& unigrams, float distortion);
};

}  // namespace tensorflow

#endif  // TENSORFLOW_CORE_KERNELS_RANGE_SAMPLER_H_