text: # Copyright 2020 InterDigital Communications, Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from compressai.ops.parametrizers import NonNegativeParametrizer
  
  
  class GDN(nn.Module):
      r"""Generalized Divisive Normalization layer.
  
      Introduced in `"Density Modeling of Images Using a Generalized Normalization
      Transformation" <https://arxiv.org/abs/1511.06281>`_,
      by Balle Johannes, Valero Laparra, and Eero P. Simoncelli, (2016).
  
      .. math::
  
         y[i] = \frac{x[i]}{\sqrt{\beta[i] + \sum_j(\gamma[j, i] * x[j]^2)}}
  
      """
  
      def __init__(self, in_channels, inverse=False, beta_min=1e-6, gamma_init=0.1):
          super().__init__()
  
          beta_min = float(beta_min)
          gamma_init = float(gamma_init)
          self.inverse = bool(inverse)
  
          self.beta_reparam = NonNegativeParametrizer(minimum=beta_min)
          beta = torch.ones(in_channels)
          beta = self.beta_reparam.init(beta)
          self.beta = nn.Parameter(beta)
  
          self.gamma_reparam = NonNegativeParametrizer()
          gamma = gamma_init * torch.eye(in_channels)
          gamma = self.gamma_reparam.init(gamma)
          self.gamma = nn.Parameter(gamma)
  
      def forward(self, x):
          _, C, _, _ = x.size()
  
          beta = self.beta_reparam(self.beta)
          gamma = self.gamma_reparam(self.gamma)
          gamma = gamma.reshape(C, C, 1, 1)
          norm = F.conv2d(x ** 2, gamma, beta)
  
          if self.inverse:
              norm = torch.sqrt(norm)
          else:
              norm = torch.rsqrt(norm)
  
          out = x * norm
  
          return out
  
  
  class GDN1(GDN):
      r"""Simplified GDN layer.
  
      Introduced in `"Computationally Efficient Neural Image Compression"
      <http://arxiv.org/abs/1912.08771>`_, by Johnston, Nick, Elad Eban, Ariel
      Gordon, and Johannes Ballé, (2019).
  
      .. math::
  
          y[i] = \frac{x[i]}{\beta[i] + \sum_j(\gamma[j, i] * |x[j]|}
  
      """
  
      def forward(self, x):
          _, C, _, _ = x.size()
  
          beta = self.beta_reparam(self.beta)
          gamma = self.gamma_reparam(self.gamma)
          gamma = gamma.reshape(C, C, 1, 1)
          norm = F.conv2d(torch.abs(x), gamma, beta)
  
          if not self.inverse:
              norm = 1.0 / norm
  
          out = x * norm
  
          return out
  