log_survival_function(x: torch.Tensor, distribution: torch.distributions.Distribution) -> torch.Tensor

ndtr(x: torch.Tensor) -> torch.Tensor

pmf_to_quantized_cdf(pmf: torch.Tensor, precision: float) -> torch.Tensor

quantization_offset(distribution: torch.distributions.Distribution) -> torch.Tensor

x_hat = torch.round(x - offset) + offset

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)

from torch import Tensor
from torch.nn import ConvTranspose2d, Module, Sequential

from ._generalized_divisive_normalization import GeneralizedDivisiveNormalization


class SynthesisTransformation(Module):
    def __init__(self, m: int, n: int):
        super(SynthesisTransformation, self).__init__()

        self._modules = Sequential(
            ConvTranspose2d(m, n, (5,), (2,), (2,), (1,)),
            GeneralizedDivisiveNormalization(n, inverse=True),
            ConvTranspose2d(n, n, (5,), (2,), (2,), (1,)),
            GeneralizedDivisiveNormalization(n, inverse=True),
            ConvTranspose2d(n, n, (5,), (2,), (2,), (1,)),
            GeneralizedDivisiveNormalization(n, inverse=True),
            ConvTranspose2d(n, 3, (5,), (2,), (2,), (1,)),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self._modules(x)

"""
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 functools
from typing import Callable, Optional, Union

import torch
import torch.nn.functional
from torch import Tensor
from torch.nn import Module, Parameter

from ._non_negative_parameterization import NonNegativeParameterization


class GeneralizedDivisiveNormalization(Module):
    """Generalized divisive normalization

    Implements an activation function that is a multivariate generalization of
    the following sigmoid-like function:

        y[i] = x[i] / (beta[i] + sum_j(gamma[j, i] * |x[j]|^alpha))^epsilon

    where ``i`` and ``j`` map over channels. This implementation never sums
    across spatial dimensions. It is similar to local response normalization,
    but much more flexible, as ``alpha``, ``beta``, ``gamma``, and ``epsilon``
    are trainable parameters.

    The method was originally described in:

        | “Density Modeling of Images using a Generalized Normalization
            Transformation”
        | Johannes Ballé, Valero Laparra, Eero P. Simoncelli
        | https://arxiv.org/abs/1511.06281

    and expanded in:

        | “End-to-end Optimized Image Compression”
        | Johannes Ballé, Valero Laparra, Eero P. Simoncelli
        | https://arxiv.org/abs/1611.01704

    Args:
        channels: number of channels in the input.
        inverse: compute the generalized divisive normalization response. If
            ``True``, compute the inverse generalized divisive normalization
            response (one step of fixed point iteration to invert the
            generalized divisive normalization; the division is replaced by
            multiplication).
        rectify: If ``True``, apply a ``ReLU`` non-linearity to the inputs
            before calculating the generalized divisive normalization response.
        alpha_parameter: A ``Tensor`` means that the value of ``alpha`` is
            fixed. A ``Callable`` can be used to determine the value of
            ``alpha`` as a function of some other parameter or tensor. This can
            be a ``Parameter``. ``None`` means that when the layer is
            initialized, a ``NonNegativeParameterization`` layer is created to
            train ``alpha`` (with a minimum value of ``1``). The default is a
            fixed value of ``1``.
        beta_parameter: A ``Tensor`` means that the value of ``beta`` is fixed.
            A ``Callable`` can be used to determine the value of ``beta`` as a
            function of some other parameter or tensor. This can be a
            ``Parameter``. ``None`` means that when the layer is initialized, a
            ``NonNegativeParameterization`` layer is created to train ``beta``
            (with a minimum value of ``1e-6``).
        epsilon_parameter: A ``Tensor`` means that the value of ``epsilon`` is
            fixed. A ``Callable`` can be used to determine the value of
            ``epsilon`` as a function of some other parameter or tensor. This
            can be a ``Parameter``. ``None`` means that when the layer is
            initialized, a ``NonNegativeParameterization`` layer is created to
            train ``epsilon`` (with a minimum value of 1e-6). The default is a
            fixed value of ``1``.
        gamma_parameter: A ``Tensor`` means that the value of ``gamma`` is
            fixed. A ``Callable`` can be used to determine the value of
            ``gamma`` as a function of some other parameter or tensor. This can
            be a ``Parameter``. ``None`` means that when the layer is
            initialized, a ``NonNegativeParameterization`` layer is created to
            train ``gamma``.
        alpha_initializer: initializes the ``alpha`` parameter. Only used if
            ``alpha`` is trained. Defaults to ``1``.
        beta_initializer: initializes the ``beta`` parameter. Only used if
            ``beta`` is created when initializing the layer. Defaults to ``1``.
        epsilon_initializer: initializes the ``epsilon`` parameter. Only used
            if ``epsilon`` is trained. Defaults to ``1``.
        gamma_initializer: initializes the ``gamma`` parameter. Only used if
            ``gamma`` is created when initializing the layer. Defaults to the
            identity multiplied by ``0.1``. A good default value for the
            diagonal is somewhere between ``0`` and ``0.5``. If set to ``0``
            and ``beta`` is initialized as ``1``, the layer is effectively
            initialized to the identity operation.
    """

    def __init__(
        self,
        channels: Union[int, Tensor],
        inverse: bool = False,
        rectify: bool = False,
        alpha_parameter: Union[float, int, Tensor, Parameter] = None,
        beta_parameter: Union[float, int, Tensor, Parameter] = None,
        epsilon_parameter: Union[float, int, Tensor, Parameter] = None,
        gamma_parameter: Union[float, int, Tensor, Parameter] = None,
        alpha_initializer: Optional[Callable[[Tensor], Tensor]] = None,
        beta_initializer: Optional[Callable[[Tensor], Tensor]] = None,
        epsilon_initializer: Optional[Callable[[Tensor], Tensor]] = None,
        gamma_initializer: Optional[Callable[[Tensor], Tensor]] = None,
    ):
        super(GeneralizedDivisiveNormalization, self).__init__()

        self._inverse = inverse

        self._rectify = rectify

        if alpha_parameter is None:
            if alpha_initializer is None:
                alpha_initializer = torch.ones

            self._reparameterized_alpha = NonNegativeParameterization(
                alpha_initializer(channels),
                minimum=1,
            )

            self._alpha_parameter = Parameter(
                self._reparameterized_alpha.initial_value,
            )
        else:
            self._alpha_parameter = alpha_parameter

        if beta_parameter is None:
            if beta_initializer is None:
                beta_initializer = torch.ones

            self._reparameterized_beta = NonNegativeParameterization(
                beta_initializer(channels),
                minimum=1e-6,
            )

            self._beta_parameter = Parameter(
                self._reparameterized_beta.initial_value,
            )
        else:
            self._beta_parameter = beta_parameter

        if epsilon_parameter is None:
            if epsilon_initializer is None:
                epsilon_initializer = torch.ones

            self._reparameterized_epsilon = NonNegativeParameterization(
                epsilon_initializer(channels),
                minimum=1e-6,
            )

            self._epsilon_parameter = Parameter(
                self._reparameterized_epsilon.initial_value,
            )
        else:
            self._epsilon_parameter = epsilon_parameter

        if gamma_parameter is None:
            if gamma_initializer is None:
                gamma_initializer = functools.partial(lambda x: 0.1 * torch.eye(x))

            self._reparameterized_gamma = NonNegativeParameterization(
                gamma_initializer(channels),
                minimum=0,
            )

            self._gamma_parameter = Parameter(
                self._reparameterized_gamma.initial_value,
            )
        else:
            self._gamma_parameter = gamma_parameter

    def forward(self, x: Tensor) -> Tensor:
        _, channels, _, _ = x.size()

        if self._rectify:
            x = torch.nn.functional.relu(x)

        y = torch.nn.functional.conv2d(
            x ** 2,
            torch.reshape(
                self._reparameterized_gamma(self._gamma_parameter),
                (channels, channels, 1, 1),
            ),
            self._reparameterized_beta(self._beta_parameter),
        )

        if self._inverse:
            return x * torch.sqrt(y)

        return x * torch.rsqrt(y)

"""
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
import torch.testing

from neuralcompression.layers import GeneralizedDivisiveNormalization


class TestGeneralizedDivisiveNormalization:
    def test_backward(self):
        x = torch.rand((1, 32, 16, 16), requires_grad=True)

        normalization = GeneralizedDivisiveNormalization(32)

        normalized = normalization(x)

        normalized.backward(x)

        assert normalized.shape == x.shape

        assert x.grad is not None

        assert x.grad.shape == x.shape

        torch.testing.assert_allclose(
            x / torch.sqrt(1 + 0.1 * (x ** 2)),
            normalized,
        )

        normalization = GeneralizedDivisiveNormalization(
            32,
            inverse=True,
        )

        normalized = normalization(x)

        normalized.backward(x)

        assert normalized.shape == x.shape

        assert x.grad is not None

        assert x.grad.shape == x.shape

        torch.testing.assert_allclose(
            x * torch.sqrt(1 + 0.1 * (x ** 2)),
            normalized,
        )

survival_function(x: torch.Tensor, distribution: torch.distributions.Distribution) -> torch.Tensor

from torch import Tensor
from torch.nn import Conv2d, Module, ReLU, Sequential


class HyperAnalysisTransformation(Module):
    def __init__(self, m: int, n: int):
        super(HyperAnalysisTransformation, self).__init__()

        self._modules = Sequential(
            Conv2d(m, n, (3,), (1,), 2),
            ReLU(inplace=True),
            Conv2d(n, n, (5,), (2,), 2),
            ReLU(inplace=True),
            Conv2d(n, n, (5,), (2,), 2),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self._modules(x)

log_cdf(x: torch.Tensor, distribution: torch.distributions.Distribution) -> torch.Tensor

log_ndtr(x: torch.Tensor) -> torch.Tensor

import torch
from torch import Tensor
from torch.autograd import Function


class _LowerBound(Function):
    @staticmethod
    def forward(
        ctx,
        tensor: Tensor,
        lower_bound: int,
        gradient: str = "identity_if_towards",
    ) -> Tensor:
        if gradient in ("disconnected", "identity", "identity_if_towards"):
            ctx.gradient = gradient
        else:
            raise ValueError

        ctx.mask = tensor.ge(lower_bound)

        return torch.clamp(tensor, lower_bound)

    @staticmethod
    def backward(ctx, grad_output: Tensor):
        if ctx.gradient == "identity_if_towards":
            grad_output *= torch.logical_or(ctx.mask, grad_output.lt(0.0))

        if ctx.gradient == "disconnected":
            grad_output *= ctx.mask

        return grad_output.type(grad_output.dtype), None


lower_bound = _LowerBound.apply

lower_tail(distribution: torch.distributions.Distribution, tail_mass: float) -> float

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)

upper_tail(distribution: torch.distributions.Distribution, tail_mass: float) -> float