Repository: facebookresearch/NeuralCompression (language: python)

Issue:
`pmf_to_quantized_cdf` op
Transforms a probability mass function (PMF) into a quantized cumulative distribution function (CDF) for entropy coding.

```Python
pmf_to_quantized_cdf(pmf: torch.Tensor, precision: float) -> torch.Tensor
```


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index b41f241f..c9dfb1dc 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -12,6 +12,7 @@ neuralcompression.functional
 .. autofunction:: dense_image_warp
 .. autofunction:: quantization_offset
 .. autofunction:: ndtr
+.. autofunction:: pmf_to_quantized_cdf
 .. autofunction:: survival_function
 .. autofunction:: log_ndtr
 .. autofunction:: log_survival_function
diff --git a/mypy.ini b/mypy.ini
index 572a0a18..c1c1086f 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -1,5 +1,5 @@
 [mypy]
 
 # modules that don't play well with mypy
-[mypy-torchvision.*,PIL.*,compressai.*,tqdm.*,lpips.*,fvcore.*]
+[mypy-torchvision.*,PIL.*,compressai.*,tqdm.*,lpips.*,fvcore.*,neuralcompression.ext.*]
 ignore_missing_imports=True
diff --git a/neuralcompression/ext/__init__.py b/neuralcompression/ext/__init__.py
new file mode 100644
index 00000000..d5b4a146
--- /dev/null
+++ b/neuralcompression/ext/__init__.py
@@ -0,0 +1,1 @@
+from ._pmf_to_quantized_cdf import pmf_to_quantized_cdf
diff --git a/neuralcompression/ext/pmf_to_quantized_cdf_py.cc b/neuralcompression/ext/pmf_to_quantized_cdf_py.cc
new file mode 100644
index 00000000..638cb62b
--- /dev/null
+++ b/neuralcompression/ext/pmf_to_quantized_cdf_py.cc
@@ -0,0 +1,77 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <cmath>
+#include <numeric>
+#include <string>
+#include <vector>
+
+std::vector<uint32_t> pmf_to_quantized_cdf(const std::vector<float> &pmf,
+                                           int precision) {
+  std::vector<uint32_t> cdf(pmf.size() + 1);
+
+  cdf[0] = 0;
+
+  std::transform(pmf.begin(), pmf.end(), cdf.begin() + 1,
+                 [=](float p) { return std::round(p * (1 << precision)); });
+
+  const uint32_t total = std::accumulate(cdf.begin(), cdf.end(), 0);
+
+  if (total == 0)
+    throw std::domain_error("");
+
+  std::transform(cdf.begin(), cdf.end(), cdf.begin(),
+                 [precision, total](uint32_t p) {
+                   return ((static_cast<uint64_t>(1 << precision) * p) / total);
+                 });
+
+  std::partial_sum(cdf.begin(), cdf.end(), cdf.begin());
+
+  cdf.back() = 1 << precision;
+
+  for (int i = 0; i < static_cast<int>(cdf.size() - 1); ++i) {
+    if (cdf[i] == cdf[i + 1]) {
+      uint32_t peak = ~0u;
+
+      int best = -1;
+
+      for (int j = 0; j < static_cast<int>(cdf.size()) - 1; ++j) {
+        uint32_t frequency = cdf[j + 1] - cdf[j];
+
+        if (frequency > 1 && frequency < peak) {
+          peak = frequency;
+          best = j;
+        }
+      }
+
+      assert(best != -1);
+
+      if (best < i) {
+        for (int j = best + 1; j <= i; ++j) {
+          cdf[j]--;
+        }
+      } else {
+        assert(best > i);
+
+        for (int j = i + 1; j <= best; ++j) {
+          cdf[j]++;
+        }
+      }
+    }
+  }
+
+  assert(cdf[0] == 0);
+
+  assert(cdf.back() == (1 << precision));
+
+  for (int i = 0; i < static_cast<int>(cdf.size()) - 1; ++i) {
+    assert(cdf[i + 1] > cdf[i]);
+  }
+
+  return cdf;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("pmf_to_quantized_cdf", &pmf_to_quantized_cdf, "");
+}
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 3c8d8249..361db3a6 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -13,6 +13,7 @@
 from ._lower_bound import lower_bound
 from ._lower_tail import lower_tail
 from ._ndtr import ndtr
+from ._pmf_to_quantized_cdf import pmf_to_quantized_cdf
 from ._quantization_offset import quantization_offset
 from ._soft_round import (
     soft_round,
diff --git a/neuralcompression/functional/_pmf_to_quantized_cdf.py b/neuralcompression/functional/_pmf_to_quantized_cdf.py
new file mode 100644
index 00000000..e23807c4
--- /dev/null
+++ b/neuralcompression/functional/_pmf_to_quantized_cdf.py
@@ -0,0 +1,61 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from torch import Tensor
+
+import neuralcompression.ext
+
+
+def pmf_to_quantized_cdf(pmf: Tensor, precision: int) -> Tensor:
+    """Transforms a probability mass function (PMF) into a quantized cumulative
+    distribution function (CDF) for entropy coding.
+
+    Because this op uses floating-point operations 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 transformed, saved, and used on the multiple platforms.
+
+    After quantization, if the PMF does not sum to :math:`2^{precision}`, then
+    some values of PMF are increased or decreased until the sum is equal to
+    :math:`2^{precision}`.
+
+    Args:
+        pmf: a probability mass function (PMF).
+            The PMF is **not** normalized by this op. The user is responsible
+            for normalizing the PMF, if necessary.
+        precision: the number of bits for probability quantization, must be
+        greater than 0 and less than or equal to 32.
+
+    Returns:
+        the quantized CDF.
+    """
+    if torch.all(pmf <= 0.0).item():
+        message = """
+        `pmf` has a non-finite or negative member. Please check for numerical
+        problems in the probability computation."
+        """
+
+        raise ValueError(message)
+
+    if pmf.shape[-1] < 2:
+        message = """
+        `pmf` shape should be at least two in last axis.
+        """
+
+        raise ValueError(message)
+
+    if precision < 1 or precision > 32:
+        message = """
+        `precision` must be greater than 0 and less than or equal to 32.
+        """
+
+        raise ValueError(message)
+
+    cdf = neuralcompression.ext.pmf_to_quantized_cdf(pmf.tolist(), precision)
+
+    return torch.tensor(cdf)
diff --git a/pyproject.toml b/pyproject.toml
index 295c4e42..5d3b08d3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,5 +2,6 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "setuptools>=58.2.0",
+    "torch>=1.9.1",
     "wheel>=0.37.0",
 ]
diff --git a/requirements.txt b/requirements.txt
index 6f9d4e30..63207673 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
-torch==1.9.0
-torchvision==0.10.0
+torch==1.9.1
+torchvision==0.10.1
 jax==0.2.17
 jaxlib==0.1.69
 compressai==1.1.5
diff --git a/setup.py b/setup.py
index 4b309d6a..aff9f7d7 100644
--- a/setup.py
+++ b/setup.py
@@ -6,8 +6,11 @@
 """
 import os
 import re
+from pathlib import Path
 
 from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CppExtension
+
 
 # from https://github.com/facebookresearch/ClassyVision/blob/master/setup.py
 # get version string from module
@@ -31,10 +34,9 @@
     "jaxlib>=0.1.65",
     "lpips>=0.1.3",
     "pytorchvideo>=0.1.3",
-    "torch>=1.8.1",
-    "torchmetrics>=0.3.2",
+    "torch==1.9.1",
     "torchmetrics>=0.3.2",
-    "torchvision>=0.9.1",
+    "torchvision==0.10.1",
     "tqdm>=4.61.0",
 ]
 
@@ -80,4 +82,18 @@
             "sphinx>=4.2.0",
         ],
     },
+    ext_modules=[
+        CppExtension(
+            "neuralcompression.ext._pmf_to_quantized_cdf",
+            [
+                str(
+                    Path(__file__).resolve().parent
+                    / "neuralcompression"
+                    / "ext"
+                    / "pmf_to_quantized_cdf_py.cc"
+                )
+            ],
+        ),
+    ],
+    cmdclass={"build_ext": BuildExtension},
 )


Test patch:
diff --git a/tests/functional/test_pmf_to_quantized_cdf.py b/tests/functional/test_pmf_to_quantized_cdf.py
new file mode 100644
index 00000000..55a1219f
--- /dev/null
+++ b/tests/functional/test_pmf_to_quantized_cdf.py
@@ -0,0 +1,25 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import pytest
+import torch
+import torch.testing
+
+from neuralcompression.functional import pmf_to_quantized_cdf
+
+
+def test_pmf_to_quantized_cdf():
+    with pytest.raises(ValueError):
+        pmf_to_quantized_cdf(torch.tensor([0.0]), 0)
+
+    with pytest.raises(ValueError):
+        pmf_to_quantized_cdf(torch.tensor([0.0]), 17)
+
+    torch.testing.assert_close(
+        pmf_to_quantized_cdf(torch.tensor([1, 2, 3]), 8),
+        torch.tensor([0, 42, 127, 256]),
+    )