import tensorflow.compat.v1 as tf
import tensorflow_compression as tfc


class AnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(AnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_0")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_1")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_2")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_3", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=None),
        ]
        super(AnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_0", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_1", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_2", inverse=True)),
            tfc.SignalConv2D(
                3, (5, 5), name="layer_3", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=None),
        ]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_0", corr=True, strides_down=1,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
                padding="same_zeros", use_bias=False,
                activation=None)
        ]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (3, 3), name="layer_2", corr=False, strides_up=1,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=None),
        ]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


# Architecture (mean-scale, no context model) based on Table 1 of https://papers.nips.cc/paper/8275-joint-autoregressive-and-hierarchical-priors-for-learned-image-compression.pdf
class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                int(self.num_filters * 1.5), (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (3, 3), name="layer_2", corr=False, strides_up=1,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=None),
        ]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor

e3ddc5e62d165ade5b2577f9400fe97ef7f070aa
<mask token>


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='igdn_0', inverse=True)),
            tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr
            =False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_1', inverse=True)), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_2', corr=
            False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_2', inverse=True)), tfc.
            SignalConv2D(3, (5, 5), name='layer_3', corr=False, strides_up=
            2, padding='same_zeros', use_bias=True, activation=None)]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (3, 3), name=
            'layer_0', corr=True, strides_down=1, padding='same_zeros',
            use_bias=True, activation=tf.nn.relu), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_1', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tf.nn.relu),
            tfc.SignalConv2D(self.num_output_filters, (5, 5), name=
            'layer_2', corr=True, strides_down=2, padding='same_zeros',
            use_bias=False, activation=None)]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1',
            corr=False, strides_up=2, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=tf.nn.relu), tfc.
            SignalConv2D(self.num_output_filters, (3, 3), name='layer_2',
            corr=False, strides_up=1, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=None)]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(int(self.num_filters * 1.5), (5, 5), name=
            'layer_1', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_output_filters, (3, 3), name=
            'layer_2', corr=False, strides_up=1, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=None)]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor

<mask token>


class AnalysisTransform(tf.keras.layers.Layer):
    <mask token>

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(AnalysisTransform, self).__init__(*args, **kwargs)
    <mask token>
    <mask token>


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='igdn_0', inverse=True)),
            tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr
            =False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_1', inverse=True)), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_2', corr=
            False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_2', inverse=True)), tfc.
            SignalConv2D(3, (5, 5), name='layer_3', corr=False, strides_up=
            2, padding='same_zeros', use_bias=True, activation=None)]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (3, 3), name=
            'layer_0', corr=True, strides_down=1, padding='same_zeros',
            use_bias=True, activation=tf.nn.relu), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_1', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tf.nn.relu),
            tfc.SignalConv2D(self.num_output_filters, (5, 5), name=
            'layer_2', corr=True, strides_down=2, padding='same_zeros',
            use_bias=False, activation=None)]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1',
            corr=False, strides_up=2, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=tf.nn.relu), tfc.
            SignalConv2D(self.num_output_filters, (3, 3), name='layer_2',
            corr=False, strides_up=1, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=None)]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(int(self.num_filters * 1.5), (5, 5), name=
            'layer_1', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_output_filters, (3, 3), name=
            'layer_2', corr=False, strides_up=1, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=None)]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor

<mask token>


class AnalysisTransform(tf.keras.layers.Layer):
    <mask token>

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(AnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=True, strides_down=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='gdn_0')), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr=
            True, strides_down=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='gdn_1')), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_2', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tfc.GDN(name=
            'gdn_2')), tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_3', corr=True, strides_down=2, padding='same_zeros',
            use_bias=True, activation=None)]
        super(AnalysisTransform, self).build(input_shape)
    <mask token>


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='igdn_0', inverse=True)),
            tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr
            =False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_1', inverse=True)), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_2', corr=
            False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_2', inverse=True)), tfc.
            SignalConv2D(3, (5, 5), name='layer_3', corr=False, strides_up=
            2, padding='same_zeros', use_bias=True, activation=None)]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (3, 3), name=
            'layer_0', corr=True, strides_down=1, padding='same_zeros',
            use_bias=True, activation=tf.nn.relu), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_1', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tf.nn.relu),
            tfc.SignalConv2D(self.num_output_filters, (5, 5), name=
            'layer_2', corr=True, strides_down=2, padding='same_zeros',
            use_bias=False, activation=None)]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1',
            corr=False, strides_up=2, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=tf.nn.relu), tfc.
            SignalConv2D(self.num_output_filters, (3, 3), name='layer_2',
            corr=False, strides_up=1, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=None)]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(int(self.num_filters * 1.5), (5, 5), name=
            'layer_1', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_output_filters, (3, 3), name=
            'layer_2', corr=False, strides_up=1, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=None)]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor

import tensorflow.compat.v1 as tf
import tensorflow_compression as tfc


class AnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(AnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=True, strides_down=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='gdn_0')), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr=
            True, strides_down=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='gdn_1')), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_2', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tfc.GDN(name=
            'gdn_2')), tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_3', corr=True, strides_down=2, padding='same_zeros',
            use_bias=True, activation=None)]
        super(AnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, activation=tfc.GDN(name='igdn_0', inverse=True)),
            tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1', corr
            =False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_1', inverse=True)), tfc.
            SignalConv2D(self.num_filters, (5, 5), name='layer_2', corr=
            False, strides_up=2, padding='same_zeros', use_bias=True,
            activation=tfc.GDN(name='igdn_2', inverse=True)), tfc.
            SignalConv2D(3, (5, 5), name='layer_3', corr=False, strides_up=
            2, padding='same_zeros', use_bias=True, activation=None)]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (3, 3), name=
            'layer_0', corr=True, strides_down=1, padding='same_zeros',
            use_bias=True, activation=tf.nn.relu), tfc.SignalConv2D(self.
            num_filters, (5, 5), name='layer_1', corr=True, strides_down=2,
            padding='same_zeros', use_bias=True, activation=tf.nn.relu),
            tfc.SignalConv2D(self.num_output_filters, (5, 5), name=
            'layer_2', corr=True, strides_down=2, padding='same_zeros',
            use_bias=False, activation=None)]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_filters, (5, 5), name='layer_1',
            corr=False, strides_up=2, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=tf.nn.relu), tfc.
            SignalConv2D(self.num_output_filters, (3, 3), name='layer_2',
            corr=False, strides_up=1, padding='same_zeros', use_bias=True,
            kernel_parameterizer=None, activation=None)]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [tfc.SignalConv2D(self.num_filters, (5, 5), name=
            'layer_0', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(int(self.num_filters * 1.5), (5, 5), name=
            'layer_1', corr=False, strides_up=2, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=tf.nn.relu
            ), tfc.SignalConv2D(self.num_output_filters, (3, 3), name=
            'layer_2', corr=False, strides_up=1, padding='same_zeros',
            use_bias=True, kernel_parameterizer=None, activation=None)]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor

import tensorflow.compat.v1 as tf
import tensorflow_compression as tfc


class AnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(AnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_0")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_1")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_2")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_3", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=None),
        ]
        super(AnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(SynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_0", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_1", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_2", inverse=True)),
            tfc.SignalConv2D(
                3, (5, 5), name="layer_3", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=None),
        ]
        super(SynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperAnalysisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_0", corr=True, strides_down=1,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
                padding="same_zeros", use_bias=False,
                activation=None)
        ]
        super(HyperAnalysisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super(HyperSynthesisTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (3, 3), name="layer_2", corr=False, strides_up=1,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=None),
        ]
        super(HyperSynthesisTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


# Architecture (mean-scale, no context model) based on Table 1 of https://papers.nips.cc/paper/8275-joint-autoregressive-and-hierarchical-priors-for-learned-image-compression.pdf
class MBT2018HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""

    def __init__(self, num_filters, num_output_filters=None, *args, **kwargs):
        self.num_filters = num_filters
        if num_output_filters is None:  # default to the same
            num_output_filters = num_filters
        self.num_output_filters = num_output_filters
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                int(self.num_filters * 1.5), (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_output_filters, (3, 3), name="layer_2", corr=False, strides_up=1,
                padding="same_zeros", use_bias=True, kernel_parameterizer=None,
                activation=None),
        ]
        super().build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor
