Repository: facebookresearch/NeuralCompression (language: python)

Issue:
`AnalysisTransformation2D` layer
```Python
from torch import Tensor
from torch.nn import Conv2d, Module, Sequential

from ._generalized_divisive_normalization import GeneralizedDivisiveNormalization


class AnalysisTransformation(Module):
    def __init__(self, m: int, n: int):
        super(AnalysisTransformation, self).__init__()

        self._modules = Sequential(
            Conv2d(3, m, (5,), (2,), 2),
            GeneralizedDivisiveNormalization(m),
            Conv2d(m, m, (5,), (2,), 2),
            GeneralizedDivisiveNormalization(m),
            Conv2d(m, m, (5,), (2,), 2),
            GeneralizedDivisiveNormalization(m),
            Conv2d(m, n, (5,), (2,), 2),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self._modules(x)
```

Prerequisites 

- [x] #85 


Discussion / hints:

Fix patch:
diff --git a/docs/source/layers.rst b/docs/source/layers.rst
index 3e234874..d1472807 100644
--- a/docs/source/layers.rst
+++ b/docs/source/layers.rst
@@ -3,6 +3,13 @@ neuralcompression.layers
 
 .. currentmodule:: neuralcompression.layers
 
+AnalysisTransformation2D
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AnalysisTransformation
+  :members: __getitem__
+  :special-members:
+
 GeneralizedDivisiveNormalization
 ~~~~~~~~~~~~~~~~~~~~
 
diff --git a/neuralcompression/layers/__init__.py b/neuralcompression/layers/__init__.py
index 066f094f..dfab4aa2 100644
--- a/neuralcompression/layers/__init__.py
+++ b/neuralcompression/layers/__init__.py
@@ -5,6 +5,7 @@
 LICENSE file in the root directory of this source tree.
 """
 
+from ._analysis_transformation_2d import AnalysisTransformation2D
 from ._continuous_entropy import ContinuousEntropy
 from ._generalized_divisive_normalization import GeneralizedDivisiveNormalization
 from ._hyper_analysis_transformation_2d import HyperAnalysisTransformation2D
diff --git a/neuralcompression/layers/_analysis_transformation_2d.py b/neuralcompression/layers/_analysis_transformation_2d.py
new file mode 100644
index 00000000..2d251253
--- /dev/null
+++ b/neuralcompression/layers/_analysis_transformation_2d.py
@@ -0,0 +1,29 @@
+"""
+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, Module, Sequential
+
+from ._generalized_divisive_normalization import GeneralizedDivisiveNormalization
+
+
+class AnalysisTransformation2D(Module):
+    def __init__(self, m: int, n: int):
+        super(AnalysisTransformation2D, self).__init__()
+
+        self.model = Sequential(
+            Conv2d(3, m, (5, 5), (2, 2), 2),
+            GeneralizedDivisiveNormalization(m),
+            Conv2d(m, m, (5, 5), (2, 2), 2),
+            GeneralizedDivisiveNormalization(m),
+            Conv2d(m, m, (5, 5), (2, 2), 2),
+            GeneralizedDivisiveNormalization(m),
+            Conv2d(m, n, (5, 5), (2, 2), 2),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.model(x)


Test patch:
diff --git a/tests/layers/test_analysis_transformation_2d.py b/tests/layers/test_analysis_transformation_2d.py
new file mode 100644
index 00000000..89e5c557
--- /dev/null
+++ b/tests/layers/test_analysis_transformation_2d.py
@@ -0,0 +1,18 @@
+"""
+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 AnalysisTransformation2D
+
+
+class TestAnalysisTransformation2D:
+    transformation = AnalysisTransformation2D(28, 28)
+
+    x = torch.rand((16, 3, 28, 28))
+
+    assert transformation(x).shape == Size([16, 28, 2, 2])