Repository: facebookresearch/NeuralCompression (language: python)

Issue:
`HyperSynthesisTransformation` layer
```Python
from torch import Tensor
from torch.nn import Conv2d, ConvTranspose2d, Module, ReLU, Sequential


class HyperSynthesisTransformation(Module):
    def __init__(self, m: int, n: int):
        super(HyperSynthesisTransformation, self).__init__()

        self._modules = Sequential(
            ConvTranspose2d(n, n, (5, 5), (2,), (2,), (1,)),
            ReLU(inplace=True),
            ConvTranspose2d(n, n, (5, 5), (2,), (2,), (1,)),
            ReLU(inplace=True),
            Conv2d(n, m, (3, 3), (1,), 1),
            ReLU(inplace=True),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self._modules(x)
```


Discussion / hints:

Fix patch:
diff --git a/docs/source/layers.rst b/docs/source/layers.rst
index 15a6a670..202e38c9 100644
--- a/docs/source/layers.rst
+++ b/docs/source/layers.rst
@@ -3,6 +3,13 @@ neuralcompression.layers
 
 .. currentmodule:: neuralcompression.layers
 
+HyperSynthesisTransformation2D
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: HyperSynthesisTransformation2D
+  :members: __getitem__
+  :special-members:
+
 SimplifiedGDN
 ~~~~~~~~~~~~~
 
diff --git a/neuralcompression/layers/__init__.py b/neuralcompression/layers/__init__.py
index 7bd90b46..66773c02 100644
--- a/neuralcompression/layers/__init__.py
+++ b/neuralcompression/layers/__init__.py
@@ -4,4 +4,6 @@
 This source code is licensed under the MIT license found in the
 LICENSE file in the root directory of this source tree.
 """
+
+from ._hyper_synthesis_transformation_2d import HyperSynthesisTransformation2D
 from .gdn import SimplifiedGDN, SimplifiedInverseGDN
diff --git a/neuralcompression/layers/_hyper_synthesis_transformation_2d.py b/neuralcompression/layers/_hyper_synthesis_transformation_2d.py
new file mode 100644
index 00000000..8c51a442
--- /dev/null
+++ b/neuralcompression/layers/_hyper_synthesis_transformation_2d.py
@@ -0,0 +1,43 @@
+"""
+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 import Tensor
+from torch.nn import Conv2d, ConvTranspose2d, Module, ReLU, Sequential
+
+
+class HyperSynthesisTransformation2D(Module):
+    """Applies the 2D hyper synthesis transformation over an input signal.
+
+    The hyper synthesis transformation is used to infer the latent
+    representation of an input signal.
+
+    The method is described in:
+
+        | “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:
+        in_channels: number of channels in the input signal.
+        features: number of inferred latent features.
+    """
+
+    def __init__(self, in_channels: int, features: int):
+        super(HyperSynthesisTransformation2D, self).__init__()
+
+        self.model = Sequential(
+            ConvTranspose2d(features, features, (5, 5), (2, 2), (2, 2), (1, 1)),
+            ReLU(inplace=True),
+            ConvTranspose2d(features, features, (5, 5), (2, 2), (2, 2), (1, 1)),
+            ReLU(inplace=True),
+            Conv2d(features, in_channels, (3, 3), (1, 1), 1),
+            ReLU(inplace=True),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.model(x)


Test patch:
diff --git a/tests/layers/__init__.py b/tests/layers/__init__.py
new file mode 100644
index 00000000..c17674b3
--- /dev/null
+++ b/tests/layers/__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/layers/test_hyper_synthesis_transformation_2d.py b/tests/layers/test_hyper_synthesis_transformation_2d.py
new file mode 100644
index 00000000..f5beaebf
--- /dev/null
+++ b/tests/layers/test_hyper_synthesis_transformation_2d.py
@@ -0,0 +1,20 @@
+"""
+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 Size
+
+from neuralcompression.layers import HyperSynthesisTransformation2D
+
+
+class TestHyperSynthesisTransformation:
+    def test_forward(self):
+        transformation = HyperSynthesisTransformation2D(28, 28)
+
+        x = torch.rand((28, 28, 5, 5))
+
+        assert transformation(x).shape == Size([28, 28, 20, 20])