Repository: facebookresearch/NeuralCompression (language: python)

Issue:
quantization_offset op
Computes distribution-dependent quantization offset.

```Python
quantization_offset(distribution: torch.distributions.Distribution) -> torch.Tensor
```

For range coding of continuous random variables, the values need to be quantized first. Typically, it is beneficial for compression performance to align the centers of the quantization bins such that one of them coincides with the mode of the distribution. With `offset` being the mode of the distribution, for instance, this can be achieved simply by computing:

```Python
x_hat = torch.round(x - offset) + offset
```

This method tries to determine the offset in a best-effort fashion, based on which statistics the `Distribution` implements. First, a method `self._quantization_offset()` is tried. If that isn’t defined, it tries in turn: `self.mode()`, `self.quantile(.5)`, then `self.mean()`. If none of these are implemented, it falls back on quantizing to integer values (i.e., an offset of zero).

The offset is always in the range [-.5, .5] as it is assumed to be combined with a round quantizer.


Discussion / hints:

Fix patch:
diff --git a/docs/source/functional.rst b/docs/source/functional.rst
index 6b2a9f2e..89350781 100644
--- a/docs/source/functional.rst
+++ b/docs/source/functional.rst
@@ -10,4 +10,5 @@ neuralcompression.functional
 .. autofunction:: hsv2rgb
 .. autofunction:: optical_flow_to_color
 .. autofunction:: dense_image_warp
+.. autofunction:: quantization_offset
 .. autofunction:: ndtr
diff --git a/neuralcompression/functional/__init__.py b/neuralcompression/functional/__init__.py
index 231f0e4c..f2be0d45 100644
--- a/neuralcompression/functional/__init__.py
+++ b/neuralcompression/functional/__init__.py
@@ -7,6 +7,7 @@
 
 from ._estimate_tails import estimate_tails
 from ._ndtr import ndtr
+from ._quantization_offset import quantization_offset
 from ._soft_round import (
     soft_round,
     soft_round_conditional_mean,
diff --git a/neuralcompression/functional/_quantization_offset.py b/neuralcompression/functional/_quantization_offset.py
new file mode 100644
index 00000000..3107f29e
--- /dev/null
+++ b/neuralcompression/functional/_quantization_offset.py
@@ -0,0 +1,47 @@
+"""
+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
+from torch.distributions import Distribution
+
+
+def quantization_offset(distribution: Distribution) -> Tensor:
+    """Computes a distribution-dependent quantization offset.
+
+    For range coding of continuous random variables, the values need to be
+    quantized first. Typically, it is beneficial for compression performance to
+    align the centers of the quantization bins such that one of them coincides
+    with the mode of the ``distribution``. With ``offset`` being the mode of
+    the distribution, for instance, this can be achieved simply by computing:
+
+    ```
+    x_hat = torch.round(x - offset) + offset
+    ```
+
+    If `distribution.mean` is undefined, quantized integer values (i.e. an
+    offset of zero) are returned.
+
+    The ``offset`` is in the range ``[-.5, .5]`` as it is assumed the returned
+    value will be combined with a rounding quantizer.
+
+    Args:
+        distribution: an object representing a continuous-valued random
+        variable.
+
+    Returns:
+        a ``torch.Tensor`` broadcastable to shape ``distribution.batch_shape``,
+        containing the determined quantization offsets. No gradients are
+        permitted to flow through the return value.
+    """
+    try:
+        offset = distribution.mean
+    except AttributeError:
+        offset = 0
+
+    with torch.no_grad():
+        return offset - torch.round(offset)


Test patch:
diff --git a/tests/functional/test_quantization_offset.py b/tests/functional/test_quantization_offset.py
new file mode 100644
index 00000000..22fe6381
--- /dev/null
+++ b/tests/functional/test_quantization_offset.py
@@ -0,0 +1,11 @@
+from torch.distributions import Gamma, Laplace, Normal
+
+from neuralcompression.functional import quantization_offset
+
+
+def test_quantization_offset():
+    assert quantization_offset(Gamma(5.0, 1.0)) == 0.0
+
+    assert quantization_offset(Laplace(-2.0, 5.0)) == 0.0
+
+    assert quantization_offset(Normal(3.0, 5.0)) == 0.0