Repository: facebookresearch/NeuralCompression (language: python)

Issue:
NoisyNormal distribution
Normal distribution with additive identically distributed (i.i.d.) uniform noise.

Prerequisites

- [x] #70 


Discussion / hints:

Fix patch:
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
index 8b9eb134..6f6cf66d 100644
--- a/docs/source/distributions.rst
+++ b/docs/source/distributions.rst
@@ -3,6 +3,13 @@ neuralcompression.distributions
 
 .. currentmodule:: neuralcompression.distributions
 
+NoisyNormal
+~~~~~~~~~~~
+
+.. autoclass:: NoisyNormal
+  :members: __getitem__
+  :special-members:
+
 UniformNoise
 ~~~~~~~~~~~~
 
diff --git a/neuralcompression/distributions/__init__.py b/neuralcompression/distributions/__init__.py
index 01bb7367..2b0d3ae3 100644
--- a/neuralcompression/distributions/__init__.py
+++ b/neuralcompression/distributions/__init__.py
@@ -5,4 +5,5 @@
 LICENSE file in the root directory of this source tree.
 """
 
+from ._noisy_normal import NoisyNormal
 from ._uniform_noise import UniformNoise
diff --git a/neuralcompression/distributions/_noisy_normal.py b/neuralcompression/distributions/_noisy_normal.py
new file mode 100644
index 00000000..21effadf
--- /dev/null
+++ b/neuralcompression/distributions/_noisy_normal.py
@@ -0,0 +1,28 @@
+"""
+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 torch.distributions import Normal
+
+from ._uniform_noise import UniformNoise
+
+
+class NoisyNormal(UniformNoise):
+    r"""Normal distribution with additive identically distributed (i.i.d.)
+        uniform noise.
+
+    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
+    """
+
+    def __init__(self, **kwargs):
+        distribution = Normal(**kwargs)
+
+        super(NoisyNormal, self).__init__(distribution)


Test patch:
diff --git a/tests/distributions/test_noisy_normal.py b/tests/distributions/test_noisy_normal.py
new file mode 100644
index 00000000..9a2ec50f
--- /dev/null
+++ b/tests/distributions/test_noisy_normal.py
@@ -0,0 +1,14 @@
+"""
+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 torch.distributions import Normal
+
+from distributions.test_uniform_noise import TestUniformNoise
+
+
+class TestNoisyNormal(TestUniformNoise):
+    distribution = Normal(0.0, 1.0)