Repository: facebookresearch/NeuralCompression (language: python)

Issue:
UniformNoise distribution
Additive independent and identically distributed (i.i.d.) uniform noise adapter distribution.

Prerequisites:

- [x] #72 
- [x] #73 
- [x] #74 
- [x] #75 
- [x] #76 
- [x] #77 


Discussion / hints:

Fix patch:
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
new file mode 100644
index 00000000..8b9eb134
--- /dev/null
+++ b/docs/source/distributions.rst
@@ -0,0 +1,11 @@
+neuralcompression.distributions
+===============================
+
+.. currentmodule:: neuralcompression.distributions
+
+UniformNoise
+~~~~~~~~~~~~
+
+.. autoclass:: UniformNoise
+  :members: __getitem__
+  :special-members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index d9f16994..2f98165b 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -14,6 +14,7 @@ for data compression.
    :caption: Package Reference
 
    data
+   distributions
    entropy_coders
    functional
    layers
diff --git a/neuralcompression/distributions/__init__.py b/neuralcompression/distributions/__init__.py
new file mode 100644
index 00000000..01bb7367
--- /dev/null
+++ b/neuralcompression/distributions/__init__.py
@@ -0,0 +1,8 @@
+"""
+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.
+"""
+
+from ._uniform_noise import UniformNoise
diff --git a/neuralcompression/distributions/_uniform_noise.py b/neuralcompression/distributions/_uniform_noise.py
new file mode 100644
index 00000000..5551fc4f
--- /dev/null
+++ b/neuralcompression/distributions/_uniform_noise.py
@@ -0,0 +1,147 @@
+"""
+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.
+"""
+
+from typing import Any, Dict, Optional
+
+import torch
+from torch import Size, Tensor
+from torch.distributions import Distribution, Uniform
+from torch.distributions.constraints import Constraint
+
+import neuralcompression.functional as ncF
+
+
+class UniformNoise(Distribution):
+    r"""Adapts a continuous distribution via additive identically distributed
+        (i.i.d.) uniform noise.
+
+    For a ``distribution``, ``UniformNoise`` models ``distribution`` after the
+    addition of independent uniform noise. Effectively, the base density
+    function is convolved with a box kernel of width one. The resulting density
+    can be efficiently evaluated via the relation:
+
+    .. math::
+        p_{\widetilde{y}}(\widetilde{y}) =
+            c(\widetilde{y} + \frac{1}{2}) - c(\widetilde{y} - \frac{1}{2})
+
+    where :math:`p` and :math:`\widetilde{y}` are the base density and the
+    unit-width uniform density, respectively, and :math:`c` is the cumulative
+    distribution function (CDF) of :math:`p`.
+
+    The method is described in Appendix 6.2. of:
+
+        | “Variational Image Compression with a Scale Hyperprior”
+        | Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang,
+            Nick Johnston
+        | https://arxiv.org/abs/1802.01436
+
+    Args:
+        distribution: an object representing a continuous-valued random
+            variable.
+    """
+
+    def __init__(self, distribution: Distribution):
+        super(UniformNoise, self).__init__(
+            distribution.event_shape,
+            distribution.batch_shape,
+        )
+
+        self._distribution = distribution
+
+        self._uniform_distribution = Uniform(-0.5, 0.5)
+
+    @property
+    def arg_constraints(self) -> Dict[str, Constraint]:
+        return {}
+
+    @property
+    def mean(self) -> Tensor:
+        return self._distribution.mean
+
+    @property
+    def quantization_offset(self) -> Tensor:
+        return ncF.quantization_offset(self._distribution)
+
+    @property
+    def support(self) -> Optional[Any]:
+        return self._distribution.support
+
+    @property
+    def variance(self) -> Tensor:
+        return self._distribution.variance
+
+    def cdf(self, value: Tensor) -> Tensor:
+        return self._distribution.cdf(value)
+
+    def entropy(self) -> Tensor:
+        return self._distribution.entropy()
+
+    def enumerate_support(self, expand: bool = True) -> Tensor:
+        return self._distribution.enumerate_support(expand)
+
+    def expand(self, batch_shape, _instance=None) -> Distribution:
+        return self._distribution.expand(batch_shape)
+
+    def icdf(self, value: Tensor) -> Tensor:
+        return self._distribution.icdf(value)
+
+    def log_cdf(self, value: Tensor) -> Tensor:
+        return ncF.log_cdf(value, self._distribution)
+
+    def log_prob(self, value: Tensor) -> Tensor:
+        log_survival_function_positive = self.log_survival_function(value + 0.5)
+        log_survival_function_negative = self.log_survival_function(value - 0.5)
+
+        log_cdf_positive = self.log_cdf(value + 0.5)
+        log_cdf_negative = self.log_cdf(value - 0.5)
+
+        condition = log_survival_function_positive < log_cdf_positive
+
+        a = torch.where(
+            condition,
+            log_survival_function_negative,
+            log_cdf_positive,
+        )
+
+        b = torch.where(
+            condition,
+            log_survival_function_positive,
+            log_cdf_negative,
+        )
+
+        return torch.log1p(-torch.exp(b - a)) + a
+
+    def log_survival_function(self, value: Tensor) -> Tensor:
+        return ncF.log_survival_function(value, self._distribution)
+
+    def lower_tail(self, tail_mass: float) -> Tensor:
+        return ncF.lower_tail(self._distribution, tail_mass)
+
+    def prob(self, value: Tensor) -> Tensor:
+        survival_function_positive = self.survival_function(value + 0.5)
+        survival_function_negative = self.survival_function(value - 0.5)
+
+        cdf_positive = self.cdf(value + 0.5)
+        cdf_negative = self.cdf(value - 0.5)
+
+        return torch.where(
+            survival_function_positive < cdf_positive,
+            survival_function_negative - survival_function_positive,
+            cdf_positive - cdf_negative,
+        )
+
+    def rsample(self, sample_shape: Size = Size()):
+        a = self._distribution.sample(sample_shape)
+        b = self._uniform_distribution.sample(sample_shape)
+
+        return a + b
+
+    def survival_function(self, value: Tensor) -> Tensor:
+        return ncF.survival_function(value, self._distribution)
+
+    def upper_tail(self, tail_mass: float) -> Tensor:
+        return ncF.upper_tail(self._distribution, tail_mass)


Test patch:
diff --git a/tests/distributions/__init__.py b/tests/distributions/__init__.py
new file mode 100644
index 00000000..c17674b3
--- /dev/null
+++ b/tests/distributions/__init__.py
@@ -0,0 +1,6 @@
+"""
+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.
+"""
diff --git a/tests/distributions/test_uniform_noise.py b/tests/distributions/test_uniform_noise.py
new file mode 100644
index 00000000..84a237b7
--- /dev/null
+++ b/tests/distributions/test_uniform_noise.py
@@ -0,0 +1,121 @@
+"""
+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.testing
+from torch import Size
+from torch.distributions import Normal
+
+import neuralcompression.functional as ncF
+from neuralcompression.distributions import UniformNoise
+
+
+class TestUniformNoise:
+    distribution = Normal(0.0, 1.0)
+
+    generator = torch.Generator()
+
+    generator.manual_seed(0xFEEEFEEE)
+
+    x = torch.normal(0.0, 1.0, (32,), generator=generator)
+
+    def test_cdf(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).cdf(self.x),
+            self.distribution.cdf(self.x),
+        )
+
+    def test_entropy(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).entropy(),
+            self.distribution.entropy(),
+        )
+
+    def test_enumerate_support(self):
+        with pytest.raises(NotImplementedError):
+            UniformNoise(self.distribution).enumerate_support()
+
+    def test_expand(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).expand(Size([32])).batch_shape,
+            Size([32]),
+        )
+
+    def test_icdf(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).icdf(self.x).shape,
+            torch.Size([32]),
+        )
+
+    def test_log_cdf(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).log_cdf(self.x),
+            ncF.log_cdf(self.x, self.distribution),
+        )
+
+    def test_log_prob(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).log_prob(self.x).shape,
+            torch.Size([32]),
+        )
+
+    def test_log_survival_function(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).log_survival_function(self.x),
+            ncF.log_survival_function(self.x, self.distribution),
+        )
+
+    def test_lower_tail(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).lower_tail(1.0),
+            ncF.lower_tail(self.distribution, 1.0),
+        )
+
+    def test_mean(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).mean,
+            torch.tensor(0.0),
+        )
+
+    def test_prob(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).prob(self.x).shape,
+            torch.Size([32]),
+        )
+
+    def test_quantization_offset(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).quantization_offset,
+            torch.tensor(0.0),
+        )
+
+    def test_rsample(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).rsample(torch.Size((5, 4))).shape,
+            torch.Size([5, 4]),
+        )
+
+    def test_support(self):
+        assert UniformNoise(self.distribution).support == self.distribution.support
+
+    def test_survival_function(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).survival_function(self.x),
+            ncF.survival_function(self.x, self.distribution),
+        )
+
+    def test_upper_tail(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).upper_tail(1.0),
+            ncF.upper_tail(self.distribution, 1.0),
+        )
+
+    def test_variance(self):
+        torch.testing.assert_equal(
+            UniformNoise(self.distribution).variance,
+            torch.tensor(1.0),
+        )