Project import generated by Copybara.

PiperOrigin-RevId: 245129345
Change-Id: Ib979ca7a93ae98bffe571a244ec78ac40b59a7fc

Author: Googler

cc/kernels/range_coder.cc

@@ -19,15 +19,15 @@ limitations under the License.
 // a digitised message", presented to the Video & Data Recording Conference,
 // held in Southampton, July 24-27, 1979.
 //
+#include "tensorflow_compression/cc/kernels/range_coder.h"
+
 #include <limits>
 #include <string>
 
 #include "tensorflow/core/lib/gtl/array_slice.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
 
-#include "tensorflow_compression/cc/kernels/range_coder.h"
-
 namespace tensorflow_compression {
 namespace gtl = tensorflow::gtl;
 using tensorflow::int32;
@@ -36,16 +36,16 @@ using tensorflow::uint32;
 using tensorflow::uint64;
 using tensorflow::uint8;
 
-RangeEncoder::RangeEncoder(int precision) : precision_(precision) {
-  CHECK_GT(precision, 0);
-  CHECK_LE(precision, 16);
-}
+void RangeEncoder::Encode(int32 lower, int32 upper, int precision,
+                          string* sink) {
+  // Input requirement: 0 < precision < 16.
+  DCHECK_GT(precision, 0);
+  DCHECK_LE(precision, 16);
 
-void RangeEncoder::Encode(int32 lower, int32 upper, string* sink) {
   // Input requirement: 0 <= lower < upper <= 2^precision.
   DCHECK_LE(0, lower);
   DCHECK_LT(lower, upper);
-  DCHECK_LE(upper, 1 << precision_);
+  DCHECK_LE(upper, 1 << precision);
 
   // `base` and `size` represent a half-open interval [base, base + size).
   // Loop invariant: 2^16 <= size <= 2^32.
@@ -69,8 +69,8 @@ void RangeEncoder::Encode(int32 lower, int32 upper, string* sink) {
   // NOTE: The max value of `size` is 2^32 and size > 0. Therefore `size * u`
   // can be rewritten as `(size - 1) * u + u` and all the computation can be
   // done in 32-bit mode. If 32-bit multiply is faster, then rewrite.
-  const uint32 a = (size * static_cast<uint64>(lower)) >> precision_;
-  const uint32 b = ((size * static_cast<uint64>(upper)) >> precision_) - 1;
+  const uint32 a = (size * static_cast<uint64>(lower)) >> precision;
+  const uint32 b = ((size * static_cast<uint64>(upper)) >> precision) - 1;
   DCHECK_LE(a, b);
 
   // Let's confirm the RHS of a, b fit in uint32 type.
@@ -301,23 +301,21 @@ void RangeEncoder::Finalize(string* sink) {
   delay_ = 0;
 }
 
-RangeDecoder::RangeDecoder(const string& source, int precision)
-    : current_(source.begin()),
-      begin_(source.begin()),
-      end_(source.end()),
-      precision_(precision) {
-  CHECK_LE(precision, 16);
-
+RangeDecoder::RangeDecoder(const string& source)
+    : current_(source.begin()), end_(source.end()) {
   Read16BitValue();
   Read16BitValue();
 }
 
-int32 RangeDecoder::Decode(gtl::ArraySlice<int32> cdf) {
+int32 RangeDecoder::Decode(gtl::ArraySlice<int32> cdf, int precision) {
+  // Input requirement: 0 < precision < 16.
+  DCHECK_GT(precision, 0);
+  DCHECK_LE(precision, 16);
+
   const uint64 size = static_cast<uint64>(size_minus1_) + 1;
   const uint64 offset =
-      ((static_cast<uint64>(value_ - base_) + 1) << precision_) - 1;
+      ((static_cast<uint64>(value_ - base_) + 1) << precision) - 1;
 
-  // This is similar to std::lower_range() with std::less_equal as comparison.
   // After the binary search, `pv` points to the smallest number v that
   // satisfies offset < (size * v) / 2^precision.
 
@@ -333,7 +331,7 @@ int32 RangeDecoder::Decode(gtl::ArraySlice<int32> cdf) {
     const auto half = len / 2;
     const int32* mid = pv + half;
     DCHECK_GE(*mid, 0);
-    DCHECK_LE(*mid, 1 << precision_);
+    DCHECK_LE(*mid, 1 << precision);
     if (size * static_cast<uint64>(*mid) <= offset) {
       pv = mid + 1;
       len -= half + 1;
@@ -349,10 +347,10 @@ int32 RangeDecoder::Decode(gtl::ArraySlice<int32> cdf) {
   // cdf.size() - 2 instead and give up detecting this error.
   CHECK_LT(pv, cdf.data() + cdf.size());
 
-  const uint32 a = (size * static_cast<uint64>(*(pv - 1))) >> precision_;
-  const uint32 b = ((size * static_cast<uint64>(*pv)) >> precision_) - 1;
-  DCHECK_LE(a, offset >> precision_);
-  DCHECK_LE(offset >> precision_, b);
+  const uint32 a = (size * static_cast<uint64>(*(pv - 1))) >> precision;
+  const uint32 b = ((size * static_cast<uint64>(*pv)) >> precision) - 1;
+  DCHECK_LE(a, offset >> precision);
+  DCHECK_LE(offset >> precision, b);
 
   base_ += a;
   size_minus1_ = b - a;
@@ -378,5 +376,4 @@ void RangeDecoder::Read16BitValue() {
     value_ |= static_cast<uint8>(*current_++);
   }
 }
-
 }  // namespace tensorflow_compression

cc/kernels/range_coder.h

@@ -26,11 +26,7 @@ namespace tensorflow_compression {
 
 class RangeEncoder {
  public:
-  // `precision` determines the granularity of probability masses passed to
-  // Encode() function below.
-  //
-  // REQUIRES: 0 < precision <= 16.
-  explicit RangeEncoder(int precision);
+  RangeEncoder() = default;
 
   // Encodes a half-open interval [lower / 2^precision, upper / 2^precision).
   // Suppose each character to be encoded is from an integer-valued
@@ -47,7 +43,8 @@ class RangeEncoder {
   // ...
   //
   // REQUIRES: 0 <= lower < upper <= 2^precision.
-  void Encode(tensorflow::int32 lower, tensorflow::int32 upper,
+  // REQUIRES: 0 < precision <= 16.
+  void Encode(tensorflow::int32 lower, tensorflow::int32 upper, int precision,
               tensorflow::string* sink);
 
   // The encode may contain some under-determined values from previous encoding.
@@ -60,18 +57,13 @@ class RangeEncoder {
   tensorflow::uint32 size_minus1_ =
       std::numeric_limits<tensorflow::uint32>::max();
   tensorflow::uint64 delay_ = 0;
-
-  const int precision_;
 };
 
 class RangeDecoder {
  public:
   // Holds a reference to `source`. The caller has to make sure that `source`
   // outlives the decoder object.
-  //
-  // REQUIRES: `precision` must be the same as the encoder's precision.
-  // REQUIRES: 0 < precision <= 16.
-  RangeDecoder(const tensorflow::string& source, int precision);
+  explicit RangeDecoder(const tensorflow::string& source);
 
   // Decodes a character from `source` using CDF. The size of `cdf` should be
   // one more than the number of the character in the alphabet.
@@ -90,9 +82,11 @@ class RangeDecoder {
   // REQUIRES: cdf.size() > 1.
   // REQUIRES: cdf[i] <= cdf[i + 1] for i = 0, 1, ..., cdf.size() - 2.
   // REQUIRES: cdf[cdf.size() - 1] <= 2^precision.
+  // REQUIRES: 0 < precision <= 16.
   //
   // In practice the last element of `cdf` should equal to 2^precision.
-  tensorflow::int32 Decode(tensorflow::gtl::ArraySlice<tensorflow::int32> cdf);
+  tensorflow::int32 Decode(tensorflow::gtl::ArraySlice<tensorflow::int32> cdf,
+                           int precision);
 
  private:
   void Read16BitValue();
@@ -103,10 +97,7 @@ class RangeDecoder {
   tensorflow::uint32 value_ = 0;
 
   tensorflow::string::const_iterator current_;
-  const tensorflow::string::const_iterator begin_;
   const tensorflow::string::const_iterator end_;
-
-  const int precision_;
 };
 
 }  // namespace tensorflow_compression

cc/kernels/range_coder_test.cc

@@ -67,11 +67,11 @@ void RangeEncodeDecodeTest(int precision, random::SimplePhilox* gen) {
     ideal_code_length[i] = -std::log2((cdf[i + 1] - cdf[i]) / normalizer);
   }
 
-  RangeEncoder encoder(precision);
+  RangeEncoder encoder;
   string encoded;
   double ideal_length = 0.0;
   for (uint8 x : data) {
-    encoder.Encode(cdf[x], cdf[x + 1], &encoded);
+    encoder.Encode(cdf[x], cdf[x + 1], precision, &encoded);
     ideal_length += ideal_code_length[x];
   }
   encoder.Finalize(&encoded);
@@ -82,9 +82,9 @@ void RangeEncodeDecodeTest(int precision, random::SimplePhilox* gen) {
             << " (ideal compression rate " << ideal_length / (8 * data.size())
             << ")";
 
-  RangeDecoder decoder(encoded, precision);
+  RangeDecoder decoder(encoded);
   for (int i = 0; i < data.size(); ++i) {
-    const int32 decoded = decoder.Decode(cdf);
+    const int32 decoded = decoder.Decode(cdf, precision);
     ASSERT_EQ(decoded, static_cast<int32>(data[i])) << i;
   }
 }
@@ -110,12 +110,12 @@ TEST(RangeCoderTest, FinalizeState0) {
   constexpr int kPrecision = 2;
 
   string output;
-  RangeEncoder encoder(kPrecision);
-  encoder.Encode(0, 2, &output);
+  RangeEncoder encoder;
+  encoder.Encode(0, 2, kPrecision, &output);
   encoder.Finalize(&output);
 
-  RangeDecoder decoder(output, kPrecision);
-  EXPECT_EQ(decoder.Decode({0, 2, 4}), 0);
+  RangeDecoder decoder(output);
+  EXPECT_EQ(decoder.Decode({0, 2, 4}, kPrecision), 0);
 }
 
 }  // namespace

cc/kernels/range_coding_kernels.cc

@@ -243,7 +243,7 @@ class RangeEncodeOp : public OpKernel {
 
     BroadcastRange<const int16, int32, N> view{data.data(), data_shape,
                                                cdf.data(), cdf_shape};
-    RangeEncoder encoder{precision_};
+    RangeEncoder encoder;
     for (int64 linear = 0; linear < data_size; ++linear) {
       const auto pair = view.Next();
 
@@ -263,7 +263,7 @@ class RangeEncodeOp : public OpKernel {
 
       const int32 lower = cdf_slice[index];
       const int32 upper = cdf_slice[index + 1];
-      encoder.Encode(lower, upper, output);
+      encoder.Encode(lower, upper, precision_, output);
     }
 
     encoder.Finalize(output);
@@ -352,7 +352,7 @@ class RangeDecodeOp : public OpKernel {
     BroadcastRange<int16, int32, N> view{output.data(), output_shape,
                                          cdf.data(), cdf_shape};
 
-    RangeDecoder decoder{encoded, precision_};
+    RangeDecoder decoder(encoded);
 
     const int64 output_size = output.size();
     const int64 cdf_size = cdf.size();
@@ -368,7 +368,7 @@ class RangeDecodeOp : public OpKernel {
       const int32* cdf_slice = pair.second;
       DCHECK_LE(cdf_slice + chip_size, cdf.data() + cdf_size);
 
-      *data = decoder.Decode({cdf_slice, chip_size});
+      *data = decoder.Decode({cdf_slice, chip_size}, precision_);
     }
     return tensorflow::Status::OK();
   }

cc/kernels/range_coding_kernels_test.cc

@@ -469,10 +469,10 @@ TEST_F(RangeCoderOpsTest, EncoderDebug) {
 }
 
 TEST_F(RangeCoderOpsTest, DecoderDebug) {
-  RangeEncoder encoder(5);
+  RangeEncoder encoder;
 
   string encoded_string;
-  encoder.Encode(16, 18, &encoded_string);
+  encoder.Encode(16, 18, 5, &encoded_string);
   encoder.Finalize(&encoded_string);
 
   Tensor encoded(DT_STRING, {});