Added PmfToQuantizedCdf op to contrib/coder in TensorFlow.

The added op transforms probability mass functions (PMF) to quantized
cumulative distribution function (CDF), which can be used by range coder ops in
contrib/coder.

The op takes greedy approach to ensure that the post-quantization probability
masses do not sum over the maximum quantized value. The op does not make any
adjustment when the post-quantization probability masses already sum less than
the maximum value.

PiperOrigin-RevId: 192827779

1 files changed, 32 insertions, 0 deletions

diff --git a/tensorflow/contrib/coder/ops/coder_ops.cc b/tensorflow/contrib/coder/ops/coder_ops.cc
index 9056d1a696..9bb171298f 100644
--- a/tensorflow/contrib/coder/ops/coder_ops.cc
+++ b/tensorflow/contrib/coder/ops/coder_ops.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
+using shape_inference::DimensionHandle;
 using shape_inference::InferenceContext;
 using shape_inference::ShapeHandle;
 
@@ -115,5 +116,36 @@ decoded: An int32 tensor with shape equal to `shape`.
 precision: The number of bits for probability quantization. Must be <= 16, and
   must match the precision used by RangeEncode that produced `encoded`.
 )doc");
+
+REGISTER_OP("PmfToQuantizedCdf")
+    .Input("pmf: float")
+    .Output("cdf: int32")
+    .Attr("precision: int >= 1")
+    .SetShapeFn([] (InferenceContext* c) {
+      ShapeHandle in;
+      TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &in));
+      DimensionHandle last;
+      TF_RETURN_IF_ERROR(c->Add(c->Dim(in, -1), 1, &last));
+      ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->ReplaceDim(in, -1, last, &out));
+      c->set_output(0, out);
+      return Status::OK();
+    })
+    .Doc(R"doc(
+Converts PMF to quantized CDF. This op uses floating-point operations
+internally. Therefore the quantized output may not be consistent across multiple
+platforms. For entropy encoders and decoders to have the same quantized CDF on
+different platforms, the quantized CDF should be produced once and saved, then
+the saved quantized CDF should be used everywhere.
+
+After quantization, if PMF sums to less than or equal to 2^precision, then this
+is equivalent to cumsum over the last dimension. This op makes no effort to make
+the sum close to 2^precision when the sum is already <= 2^precision.
+
+After quantization, if PMF sums to greater than 2^precision, then some values of
+PMF is decreased to keep the sum no more than 2^precision.
+
+Note that the input PMF is pre-quantization.
+)doc");
 // clang-format on
 }  // namespace tensorflow