dataeval 0.61.0__py3-none-any.whl

dataeval/_internal/models/tensorflow/pixelcnn.py

@@ -0,0 +1,1106 @@
+# type: ignore
+
+"""
+Source code derived from Alibi-Detect 0.11.4
+https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
+
+Original code Copyright (c) 2023 Seldon Technologies Ltd
+Licensed under Apache Software License (Apache 2.0)
+"""
+
+import functools
+import warnings
+from typing import Optional
+
+import keras
+import numpy as np
+import tensorflow as tf
+from tensorflow_probability.python.bijectors import bijector
+from tensorflow_probability.python.distributions import (
+    categorical,
+    distribution,
+    independent,
+    logistic,
+    mixture_same_family,
+    quantized_distribution,
+    transformed_distribution,
+)
+from tensorflow_probability.python.internal import (
+    dtype_util,
+    prefer_static,
+    reparameterization,
+    tensor_util,
+    tensorshape_util,
+)
+
+__all__ = [
+    "Shift",
+]
+
+
+class WeightNorm(keras.layers.Wrapper):
+    def __init__(self, layer, data_init: bool = True, **kwargs):
+        """Layer wrapper to decouple magnitude and direction of the layer's weights.
+
+        This wrapper reparameterizes a layer by decoupling the weight's
+        magnitude and direction. This speeds up convergence by improving the
+        conditioning of the optimization problem. It has an optional data-dependent
+        initialization scheme, in which initial values of weights are set as functions
+        of the first minibatch of data. Both the weight normalization and data-
+        dependent initialization are described in [Salimans and Kingma (2016)][1].
+
+        Parameters
+        ----------
+        layer
+            A `keras.layers.Layer` instance. Supported layer types are
+            `Dense`, `Conv2D`, and `Conv2DTranspose`. Layers with multiple inputs
+            are not supported.
+        data_init
+            If `True` use data dependent variable initialization.
+        **kwargs
+            Additional keyword args passed to `keras.layers.Wrapper`.
+
+        Raises
+        ------
+        ValueError
+            If `layer` is not a `keras.layers.Layer` instance.
+        """
+        if not isinstance(layer, keras.layers.Layer):
+            raise ValueError(
+                "Please initialize `WeightNorm` layer with a `keras.layers.Layer` " f"instance. You passed: {layer}"
+            )
+
+        layer_type = type(layer).__name__
+        if layer_type not in ["Dense", "Conv2D", "Conv2DTranspose"]:
+            warnings.warn(
+                "`WeightNorm` is tested only for `Dense`, `Conv2D`, and "
+                f"`Conv2DTranspose` layers. You passed a layer of type `{layer_type}`"
+            )
+
+        super().__init__(layer, **kwargs)
+
+        self.data_init = data_init
+        self._track_trackable(layer, name="layer")
+        self.filter_axis = -2 if layer_type == "Conv2DTranspose" else -1
+
+    def _compute_weights(self):
+        """Generate weights with normalization."""
+        # Determine the axis along which to expand `g` so that `g` broadcasts to
+        # the shape of `v`.
+        new_axis = -self.filter_axis - 3
+
+        self.layer.kernel = tf.nn.l2_normalize(self.v, axis=self.kernel_norm_axes) * tf.expand_dims(self.g, new_axis)
+
+    def _init_norm(self):
+        """Set the norm of the weight vector."""
+        kernel_norm = tf.sqrt(tf.reduce_sum(tf.square(self.v), axis=self.kernel_norm_axes))
+        self.g.assign(kernel_norm)
+
+    def _data_dep_init(self, inputs):
+        """Data dependent initialization."""
+        # Normalize kernel first so that calling the layer calculates
+        # `tf.dot(v, x)/tf.norm(v)` as in (5) in ([Salimans and Kingma, 2016][1]).
+        self._compute_weights()
+
+        activation = self.layer.activation
+        self.layer.activation = None
+
+        use_bias = self.layer.bias is not None
+        if use_bias:
+            bias = self.layer.bias
+            self.layer.bias = tf.zeros_like(bias)
+
+        # Since the bias is initialized as zero, setting the activation to zero and
+        # calling the initialized layer (with normalized kernel) yields the correct
+        # computation ((5) in Salimans and Kingma (2016))
+        x_init = self.layer(inputs)
+        norm_axes_out = list(range(x_init.shape.rank - 1))
+        m_init, v_init = tf.nn.moments(x_init, norm_axes_out)
+        scale_init = 1.0 / tf.sqrt(v_init + 1e-10)
+
+        self.g.assign(self.g * scale_init)
+        if use_bias:
+            self.layer.bias = bias
+            self.layer.bias.assign(-m_init * scale_init)
+        self.layer.activation = activation
+
+    def build(self, input_shape=None):
+        """Build `Layer`.
+
+        Parameters
+        ----------
+        input_shape
+            The shape of the input to `self.layer`.
+
+        Raises
+        ------
+        ValueError
+            If `Layer` does not contain a `kernel` of weights.
+        """
+        input_shape = tf.TensorShape(input_shape).as_list()
+        input_shape[0] = None
+        self.input_spec = keras.layers.InputSpec(shape=input_shape)
+
+        if not self.layer.built:
+            self.layer.build(input_shape)
+
+            if not hasattr(self.layer, "kernel"):
+                raise ValueError("`WeightNorm` must wrap a layer that contains a `kernel` for weights")
+
+            self.kernel_norm_axes = list(range(self.layer.kernel.shape.ndims))
+            self.kernel_norm_axes.pop(self.filter_axis)
+
+            self.v = self.layer.kernel
+
+            # to avoid a duplicate `kernel` variable after `build` is called
+            self.layer.kernel = None
+            self.g = self.add_weight(
+                name="g",
+                shape=(int(self.v.shape[self.filter_axis]),),
+                initializer="ones",
+                dtype=self.v.dtype,
+                trainable=True,
+            )
+            self.initialized = self.add_weight(name="initialized", dtype=tf.bool, trainable=False)
+            self.initialized.assign(False)
+
+        super().build()
+
+    @tf.function
+    def call(self, inputs):
+        """Call `Layer`."""
+        if not self.initialized:
+            if self.data_init:
+                self._data_dep_init(inputs)
+            else:  # initialize `g` as the norm of the initialized kernel
+                self._init_norm()
+
+            self.initialized.assign(True)
+
+        self._compute_weights()
+        output = self.layer(inputs)
+        return output
+
+    def compute_output_shape(self, input_shape):
+        return tf.TensorShape(self.layer.compute_output_shape(input_shape).as_list())
+
+
+class Shift(bijector.Bijector):
+    def __init__(self, shift, validate_args=False, name="shift"):
+        """Instantiates the `Shift` bijector which computes `Y = g(X; shift) = X + shift`
+        where `shift` is a numeric `Tensor`.
+
+        Parameters
+        ----------
+        shift
+            Floating-point `Tensor`.
+        validate_args
+            Python `bool` indicating whether arguments should be checked for correctness.
+        name
+            Python `str` name given to ops managed by this object.
+        """
+        with tf.name_scope(name) as name:
+            dtype = dtype_util.common_dtype([shift], dtype_hint=tf.float32)
+            self._shift = tensor_util.convert_nonref_to_tensor(shift, dtype=dtype, name="shift")
+            super().__init__(
+                forward_min_event_ndims=0,
+                is_constant_jacobian=True,
+                dtype=dtype,
+                validate_args=validate_args,
+                name=name,
+            )
+
+    @property
+    def shift(self):
+        """The `shift` `Tensor` in `Y = X + shift`."""
+        return self._shift
+
+    @classmethod
+    def _is_increasing(cls):
+        return True
+
+    def _forward(self, x):
+        return x + self.shift
+
+    def _inverse(self, y):
+        return y - self.shift
+
+    def _forward_log_det_jacobian(self, x):
+        # is_constant_jacobian = True for this bijector, hence the
+        # `log_det_jacobian` need only be specified for a single input, as this will
+        # be tiled to match `event_ndims`.
+        return tf.zeros([], dtype=dtype_util.base_dtype(x.dtype))
+
+
+class PixelCNN(distribution.Distribution):
+    """
+    Construct Pixel CNN++ distribution.
+
+    Parameters
+    ----------
+    image_shape
+        3D `TensorShape` or tuple for the `[height, width, channels]` dimensions of the image.
+    conditional_shape
+        `TensorShape` or tuple for the shape of the conditional input, or `None` if there is no conditional input.
+    num_resnet
+        The number of layers (shown in Figure 2 of [2]) within each highest-level block of Figure 2 of [1].
+    num_hierarchies
+        The number of highest-level blocks (separated by expansions/contractions of dimensions in Figure 2 of [1].)
+    num_filters
+        The number of convolutional filters.
+    num_logistic_mix
+        Number of components in the logistic mixture distribution.
+    receptive_field_dims
+        Height and width in pixels of the receptive field of the convolutional layers above and to the left
+        of a given pixel. The width (second element of the tuple) should be odd. Figure 1 (middle) of [2]
+        shows a receptive field of (3, 5) (the row containing the current pixel is included in the height).
+        The default of (3, 3) was used to produce the results in [1].
+    dropout_p
+        The dropout probability. Should be between 0 and 1.
+    resnet_activation
+        The type of activation to use in the resnet blocks. May be 'concat_elu', 'elu', or 'relu'.
+    l2_weight
+        The L2 regularization weight.
+    use_weight_norm
+        If `True` then use weight normalization (works only in Eager mode).
+    use_data_init
+        If `True` then use data-dependent initialization (has no effect if `use_weight_norm` is `False`).
+    high
+        The maximum value of the input data (255 for an 8-bit image).
+    low
+        The minimum value of the input data.
+    dtype
+        Data type of the `Distribution`.
+    name
+        The name of the `Distribution`.
+    """
+
+    def __init__(
+        self,
+        image_shape: tuple,
+        conditional_shape: Optional[tuple] = None,
+        num_resnet: int = 5,
+        num_hierarchies: int = 3,
+        num_filters: int = 160,
+        num_logistic_mix: int = 10,
+        receptive_field_dims: tuple = (3, 3),
+        dropout_p: float = 0.5,
+        resnet_activation: str = "concat_elu",
+        l2_weight: float = 0.0,
+        use_weight_norm: bool = True,
+        use_data_init: bool = True,
+        high: int = 255,
+        low: int = 0,
+        dtype=tf.float32,
+        name: str = "PixelCNN",
+    ) -> None:
+        parameters = dict(locals())
+        with tf.name_scope(name) as name:
+            super().__init__(
+                dtype=dtype,
+                reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
+                validate_args=False,
+                allow_nan_stats=True,
+                parameters=parameters,
+                name=name,
+            )
+
+            if not tensorshape_util.is_fully_defined(image_shape):
+                raise ValueError("`image_shape` must be fully defined.")
+
+            if conditional_shape is not None and not tensorshape_util.is_fully_defined(conditional_shape):
+                raise ValueError("`conditional_shape` must be fully defined.")
+
+            if tensorshape_util.rank(image_shape) != 3:
+                raise ValueError("`image_shape` must have length 3, representing [height, width, channels] dimensions.")
+
+            self._high = tf.cast(high, self.dtype)
+            self._low = tf.cast(low, self.dtype)
+            self._num_logistic_mix = num_logistic_mix
+            self.network = _PixelCNNNetwork(
+                dropout_p=dropout_p,
+                num_resnet=num_resnet,
+                num_hierarchies=num_hierarchies,
+                num_filters=num_filters,
+                num_logistic_mix=num_logistic_mix,
+                receptive_field_dims=receptive_field_dims,
+                resnet_activation=resnet_activation,
+                l2_weight=l2_weight,
+                use_weight_norm=use_weight_norm,
+                use_data_init=use_data_init,
+                dtype=dtype,
+            )
+
+            image_input_shape = tensorshape_util.concatenate([None], image_shape)
+            if conditional_shape is None:
+                input_shape = image_input_shape
+            else:
+                conditional_input_shape = tensorshape_util.concatenate([None], conditional_shape)
+                input_shape = [image_input_shape, conditional_input_shape]
+
+            self.image_shape = image_shape
+            self.conditional_shape = conditional_shape
+            self.network.build(input_shape)
+
+    def _make_mixture_dist(self, component_logits, locs, scales, return_per_feature: bool = False):
+        """Builds a mixture of quantized logistic distributions.
+
+        Parameters
+        ----------
+        component_logits
+            4D `Tensor` of logits for the Categorical distribution
+            over Quantized Logistic mixture components. Dimensions are `[batch_size,
+            height, width, num_logistic_mix]`.
+        locs
+            4D `Tensor` of location parameters for the Quantized Logistic
+            mixture components. Dimensions are `[batch_size, height, width,
+            num_logistic_mix, num_channels]`.
+        scales
+            4D `Tensor` of location parameters for the Quantized Logistic
+            mixture components. Dimensions are `[batch_size, height, width,
+            num_logistic_mix, num_channels]`.
+        return_per_feature
+            If True, return per pixel level log prob.
+
+        Returns
+        -------
+        dist
+            A quantized logistic mixture `tfp.distribution` over the input data.
+        """
+        mixture_distribution = categorical.Categorical(logits=component_logits)
+
+        # Convert distribution parameters for pixel values in
+        # `[self._low, self._high]` for use with `QuantizedDistribution`
+        locs = self._low + 0.5 * (self._high - self._low) * (locs + 1.0)
+        scales *= 0.5 * (self._high - self._low)
+        logistic_dist = quantized_distribution.QuantizedDistribution(
+            distribution=transformed_distribution.TransformedDistribution(
+                distribution=logistic.Logistic(loc=locs, scale=scales),
+                bijector=Shift(shift=tf.cast(-0.5, self.dtype)),
+            ),
+            low=self._low,
+            high=self._high,
+        )
+
+        # mixture with logistics for the loc and scale on each pixel for each component
+        dist = mixture_same_family.MixtureSameFamily(
+            mixture_distribution=mixture_distribution,
+            components_distribution=independent.Independent(logistic_dist, reinterpreted_batch_ndims=1),
+        )
+        if return_per_feature:
+            return dist
+        else:
+            return independent.Independent(dist, reinterpreted_batch_ndims=2)
+
+    def _log_prob(self, value, conditional_input=None, training=None, return_per_feature=False):
+        """Log probability function with optional conditional input.
+
+        Calculates the log probability of a batch of data under the modeled
+        distribution (or conditional distribution, if conditional input is
+        provided).
+
+        Parameters
+        ----------
+        value
+            `Tensor` or Numpy array of image data. May have leading batch
+            dimension(s), which must broadcast to the leading batch dimensions of
+            `conditional_input`.
+        conditional_input
+            `Tensor` on which to condition the distribution (e.g.
+            class labels), or `None`. May have leading batch dimension(s), which
+            must broadcast to the leading batch dimensions of `value`.
+        training
+            `bool` or `None`. If `bool`, it controls the dropout layer,
+            where `True` implies dropout is active. If `None`, it defaults to
+            `keras.backend.learning_phase()`.
+        return_per_feature
+            `bool`. If True, return per pixel level log prob.
+
+        Returns
+        -------
+        log_prob_values: `Tensor`.
+        """
+        # Determine the batch shape of the input images
+        image_batch_shape = prefer_static.shape(value)[:-3]
+
+        # Broadcast `value` and `conditional_input` to the same batch_shape
+        if conditional_input is None:
+            image_batch_and_conditional_shape = image_batch_shape
+        else:
+            conditional_input = tf.convert_to_tensor(conditional_input)
+            conditional_input_shape = prefer_static.shape(conditional_input)
+            conditional_batch_rank = prefer_static.rank(conditional_input) - tensorshape_util.rank(
+                self.conditional_shape
+            )
+            conditional_batch_shape = conditional_input_shape[:conditional_batch_rank]
+
+            image_batch_and_conditional_shape = prefer_static.broadcast_shape(
+                image_batch_shape, conditional_batch_shape
+            )
+            conditional_input = tf.broadcast_to(
+                conditional_input,
+                prefer_static.concat([image_batch_and_conditional_shape, self.conditional_shape], axis=0),
+            )
+            value = tf.broadcast_to(
+                value,
+                prefer_static.concat([image_batch_and_conditional_shape, self.event_shape], axis=0),
+            )
+
+            # Flatten batch dimension for input to Keras model
+            conditional_input = tf.reshape(
+                conditional_input,
+                prefer_static.concat([(-1,), self.conditional_shape], axis=0),
+            )
+
+        value = tf.reshape(value, prefer_static.concat([(-1,), self.event_shape], axis=0))
+
+        transformed_value = (2.0 * (value - self._low) / (self._high - self._low)) - 1.0
+        inputs = transformed_value if conditional_input is None else [transformed_value, conditional_input]
+
+        params = self.network(inputs, training=training)
+
+        num_channels = self.event_shape[-1]
+        if num_channels == 1:
+            component_logits, locs, scales = params
+        else:
+            # If there is more than one channel, we create a linear autoregressive
+            # dependency among the location parameters of the channels of a single
+            # pixel (the scale parameters within a pixel are independent). For a pixel
+            # with R/G/B channels, the `r`, `g`, and `b` saturation values are
+            # distributed as:
+            #
+            # r ~ Logistic(loc_r, scale_r)
+            # g ~ Logistic(coef_rg * r + loc_g, scale_g)
+            # b ~ Logistic(coef_rb * r + coef_gb * g + loc_b, scale_b)
+            # on the coefficients instead of split/multiply/concat
+            component_logits, locs, scales, coeffs = params
+            num_coeffs = num_channels * (num_channels - 1) // 2
+            loc_tensors = tf.split(locs, num_channels, axis=-1)
+            coef_tensors = tf.split(coeffs, num_coeffs, axis=-1)
+            channel_tensors = tf.split(value, num_channels, axis=-1)
+
+            coef_count = 0
+            for i in range(num_channels):
+                channel_tensors[i] = channel_tensors[i][..., tf.newaxis, :]
+                for j in range(i):
+                    loc_tensors[i] += channel_tensors[j] * coef_tensors[coef_count]
+                    coef_count += 1
+            locs = tf.concat(loc_tensors, axis=-1)
+
+        dist = self._make_mixture_dist(component_logits, locs, scales, return_per_feature=return_per_feature)
+        log_px = dist.log_prob(value)
+        if return_per_feature:
+            return log_px
+        else:
+            return tf.reshape(log_px, image_batch_and_conditional_shape)
+
+    def _sample_n(self, n, seed=None, conditional_input=None, training=False):
+        """Samples from the distribution, with optional conditional input.
+
+        Parameters
+        ----------
+        n
+            `int`, number of samples desired.
+        seed
+            `int`, seed for RNG. Setting a random seed enforces reproducibility
+            of the samples between sessions (not within a single session).
+        conditional_input
+            `Tensor` on which to condition the distribution (e.g.
+            class labels), or `None`.
+        training
+            `bool` or `None`. If `bool`, it controls the dropout layer,
+            where `True` implies dropout is active. If `None`, it defers to Keras'
+            handling of train/eval status.
+
+        Returns
+        -------
+        samples
+            a `Tensor` of shape `[n, height, width, num_channels]`.
+        """
+        if conditional_input is not None:
+            conditional_input = tf.convert_to_tensor(conditional_input, dtype=self.dtype)
+            conditional_event_rank = tensorshape_util.rank(self.conditional_shape)
+            conditional_input_shape = prefer_static.shape(conditional_input)
+            conditional_sample_rank = prefer_static.rank(conditional_input) - conditional_event_rank
+
+            # If `conditional_input` has no sample dimensions, prepend a sample
+            # dimension
+            if conditional_sample_rank == 0:
+                conditional_input = conditional_input[tf.newaxis, ...]
+                conditional_sample_rank = 1
+
+            # Assert that the conditional event shape in the `PixelCnnNetwork` is the
+            # same as that implied by `conditional_input`.
+            conditional_event_shape = conditional_input_shape[conditional_sample_rank:]
+            with tf.control_dependencies([tf.assert_equal(self.conditional_shape, conditional_event_shape)]):
+                conditional_sample_shape = conditional_input_shape[:conditional_sample_rank]
+                repeat = n // prefer_static.reduce_prod(conditional_sample_shape)
+                h = tf.reshape(
+                    conditional_input,
+                    prefer_static.concat([(-1,), self.conditional_shape], axis=0),
+                )
+                h = tf.tile(
+                    h,
+                    prefer_static.pad(
+                        [repeat],
+                        paddings=[[0, conditional_event_rank]],
+                        constant_values=1,
+                    ),
+                )
+
+        samples_0 = tf.random.uniform(
+            prefer_static.concat([(n,), self.event_shape], axis=0),
+            minval=-1.0,
+            maxval=1.0,
+            dtype=self.dtype,
+            seed=seed,
+        )
+        inputs = samples_0 if conditional_input is None else [samples_0, h]
+        params_0 = self.network(inputs, training=training)
+        samples_0 = self._sample_channels(*params_0, seed=seed)
+
+        image_height, image_width, _ = tensorshape_util.as_list(self.event_shape)
+
+        def loop_body(index, samples):
+            """Loop for iterative pixel sampling.
+
+            Parameters
+            ----------
+            index
+                0D `Tensor` of type `int32`. Index of the current pixel.
+            samples
+              4D `Tensor`. Images with pixels sampled in raster order, up to
+              pixel `[index]`, with dimensions `[batch_size, height, width,
+              num_channels]`.
+
+            Returns
+            -------
+            samples
+                4D `Tensor`. Images with pixels sampled in raster order, up to \
+                and including pixel `[index]`, with dimensions `[batch_size, height, \
+                width, num_channels]`.
+            """
+            inputs = samples if conditional_input is None else [samples, h]
+            params = self.network(inputs, training=training)
+            samples_new = self._sample_channels(*params, seed=seed)
+
+            # Update the current pixel
+            samples = tf.transpose(samples, [1, 2, 3, 0])
+            samples_new = tf.transpose(samples_new, [1, 2, 3, 0])
+            row, col = index // image_width, index % image_width
+            updates = samples_new[row, col, ...][tf.newaxis, ...]
+            samples = tf.tensor_scatter_nd_update(samples, [[row, col]], updates)
+            samples = tf.transpose(samples, [3, 0, 1, 2])
+
+            return index + 1, samples
+
+        index0 = tf.zeros([], dtype=tf.int32)
+
+        # Construct the while loop for sampling
+        total_pixels = image_height * image_width
+        loop_cond = lambda ind, _: tf.less(ind, total_pixels)  # noqa: E731
+        init_vars = (index0, samples_0)
+        _, samples = tf.while_loop(loop_cond, loop_body, init_vars, parallel_iterations=1)
+
+        transformed_samples = self._low + 0.5 * (self._high - self._low) * (samples + 1.0)
+        return tf.round(transformed_samples)
+
+    def _sample_channels(self, component_logits, locs, scales, coeffs=None, seed=None):
+        """Sample a single pixel-iteration and apply channel conditioning.
+
+        Parameters
+        ----------
+        component_logits
+            4D `Tensor` of logits for the Categorical distribution
+            over Quantized Logistic mixture components. Dimensions are `[batch_size,
+            height, width, num_logistic_mix]`.
+        locs
+            4D `Tensor` of location parameters for the Quantized Logistic
+            mixture components. Dimensions are `[batch_size, height, width,
+            num_logistic_mix, num_channels]`.
+        scales
+            4D `Tensor` of location parameters for the Quantized Logistic
+            mixture components. Dimensions are `[batch_size, height, width,
+            num_logistic_mix, num_channels]`.
+        coeffs
+            4D `Tensor` of coefficients for the linear dependence among color
+            channels, or `None` if there is only one channel. Dimensions are
+            `[batch_size, height, width, num_logistic_mix, num_coeffs]`, where
+            `num_coeffs = num_channels * (num_channels - 1) // 2`.
+        seed
+            `int`, random seed.
+
+        Returns
+        -------
+        samples
+            4D `Tensor` of sampled image data with autoregression among \
+            channels. Dimensions are `[batch_size, height, width, num_channels]`.
+        """
+        num_channels = self.event_shape[-1]
+
+        # sample mixture components once for the entire pixel
+        component_dist = categorical.Categorical(logits=component_logits)
+        mask = tf.one_hot(indices=component_dist.sample(seed=seed), depth=self._num_logistic_mix)
+        mask = tf.cast(mask[..., tf.newaxis], self.dtype)
+
+        # apply mixture component mask and separate out RGB parameters
+        masked_locs = tf.reduce_sum(locs * mask, axis=-2)
+        loc_tensors = tf.split(masked_locs, num_channels, axis=-1)
+        masked_scales = tf.reduce_sum(scales * mask, axis=-2)
+        scale_tensors = tf.split(masked_scales, num_channels, axis=-1)
+
+        if coeffs is not None:
+            num_coeffs = num_channels * (num_channels - 1) // 2
+            masked_coeffs = tf.reduce_sum(coeffs * mask, axis=-2)
+            coef_tensors = tf.split(masked_coeffs, num_coeffs, axis=-1)
+
+        channel_samples = []
+        coef_count = 0
+        for i in range(num_channels):
+            loc = loc_tensors[i]
+            for c in channel_samples:
+                loc += c * coef_tensors[coef_count]
+                coef_count += 1
+
+            logistic_samp = logistic.Logistic(loc=loc, scale=scale_tensors[i]).sample(seed=seed)
+            logistic_samp = tf.clip_by_value(logistic_samp, -1.0, 1.0)
+            channel_samples.append(logistic_samp)
+
+        return tf.concat(channel_samples, axis=-1)
+
+    def _batch_shape(self):
+        return tf.TensorShape([])
+
+    def _event_shape(self):
+        return tf.TensorShape(self.image_shape)
+
+
+class _PixelCNNNetwork(keras.layers.Layer):
+    """Keras `Layer` to parameterize a Pixel CNN++ distribution.
+    This is a Keras implementation of the Pixel CNN++ network, as described in
+    Salimans et al. (2017)[1] and van den Oord et al. (2016)[2].
+    (https://github.com/openai/pixel-cnn).
+    #### References
+    [1]: Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma.
+       PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture
+       Likelihood and Other Modifications. In _International Conference on
+       Learning Representations_, 2017.
+       https://pdfs.semanticscholar.org/9e90/6792f67cbdda7b7777b69284a81044857656.pdf
+       Additional details at https://github.com/openai/pixel-cnn
+    [2]: Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt,
+       Alex Graves, and Koray Kavukcuoglu. Conditional Image Generation with
+       PixelCNN Decoders. In _30th Conference on Neural Information Processing
+       Systems_, 2016.
+       https://papers.nips.cc/paper/6527-conditional-image-generation-with-pixelcnn-decoders.pdf.
+    """
+
+    def __init__(
+        self,
+        dropout_p: float = 0.5,
+        num_resnet: int = 5,
+        num_hierarchies: int = 3,
+        num_filters: int = 160,
+        num_logistic_mix: int = 10,
+        receptive_field_dims: tuple = (3, 3),
+        resnet_activation: str = "concat_elu",
+        l2_weight: float = 0.0,
+        use_weight_norm: bool = True,
+        use_data_init: bool = True,
+        dtype=tf.float32,
+    ) -> None:
+        """Initialize the neural network for the Pixel CNN++ distribution.
+
+        Parameters
+        ----------
+        dropout_p
+            `float`, the dropout probability. Should be between 0 and 1.
+        num_resnet
+            `int`, the number of layers (shown in Figure 2 of [2]) within
+            each highest-level block of Figure 2 of [1].
+        num_hierarchies
+            `int`, the number of hightest-level blocks (separated by
+            expansions/contractions of dimensions in Figure 2 of [1].)
+        num_filters
+            `int`, the number of convolutional filters.
+        num_logistic_mix
+            `int`, number of components in the logistic mixture
+            distribution.
+        receptive_field_dims
+            `tuple`, height and width in pixels of the receptive
+            field of the convolutional layers above and to the left of a given
+            pixel. The width (second element of the tuple) should be odd. Figure 1
+            (middle) of [2] shows a receptive field of (3, 5) (the row containing
+            the current pixel is included in the height). The default of (3, 3) was
+            used to produce the results in [1].
+        resnet_activation
+            `string`, the type of activation to use in the resnet
+            blocks. May be 'concat_elu', 'elu', or 'relu'.
+        l2_weight
+            `float`, the L2 regularization weight.
+        use_weight_norm
+            `bool`, if `True` then use weight normalization.
+        use_data_init
+            `bool`, if `True` then use data-dependent initialization
+            (has no effect if `use_weight_norm` is `False`).
+        dtype
+            Data type of the layer.
+        """
+        super().__init__(dtype=dtype)
+        self._dropout_p = dropout_p
+        self._num_resnet = num_resnet
+        self._num_hierarchies = num_hierarchies
+        self._num_filters = num_filters
+        self._num_logistic_mix = num_logistic_mix
+        self._receptive_field_dims = receptive_field_dims  # first set desired receptive field, then infer kernel
+        self._resnet_activation = resnet_activation
+        self._l2_weight = l2_weight
+
+        if use_weight_norm:
+
+            def layer_wrapper(layer):
+                def wrapped_layer(*args, **kwargs):
+                    return WeightNorm(layer(*args, **kwargs), data_init=use_data_init)
+
+                return wrapped_layer
+
+            self._layer_wrapper = layer_wrapper
+        else:
+            self._layer_wrapper = lambda layer: layer
+
+    def build(self, input_shape):
+        dtype = self.dtype
+        if len(input_shape) == 2:
+            batch_image_shape, batch_conditional_shape = input_shape
+            conditional_input = keras.layers.Input(shape=batch_conditional_shape[1:], dtype=dtype)
+        else:
+            batch_image_shape = input_shape
+            conditional_input = None
+
+        image_shape = batch_image_shape[1:]
+        image_input = keras.layers.Input(shape=image_shape, dtype=dtype)
+
+        if self._resnet_activation == "concat_elu":
+            activation = keras.layers.Lambda(lambda x: tf.nn.elu(tf.concat([x, -x], axis=-1)), dtype=dtype)
+        else:
+            activation = keras.activations.get(self._resnet_activation)
+
+        # Define layers with default inputs and layer wrapper applied
+        Conv2D = functools.partial(  # pylint:disable=invalid-name
+            self._layer_wrapper(keras.layers.Convolution2D),
+            filters=self._num_filters,
+            padding="same",
+            kernel_regularizer=keras.regularizers.l2(self._l2_weight),
+            dtype=dtype,
+        )
+
+        Dense = functools.partial(  # pylint:disable=invalid-name
+            self._layer_wrapper(keras.layers.Dense),
+            kernel_regularizer=keras.regularizers.l2(self._l2_weight),
+            dtype=dtype,
+        )
+
+        Conv2DTranspose = functools.partial(  # pylint:disable=invalid-name
+            self._layer_wrapper(keras.layers.Conv2DTranspose),
+            filters=self._num_filters,
+            padding="same",
+            strides=(2, 2),
+            kernel_regularizer=keras.regularizers.l2(self._l2_weight),
+            dtype=dtype,
+        )
+
+        rows, cols = self._receptive_field_dims
+
+        # Define the dimensions of the valid (unmasked) areas of the layer kernels
+        # for stride 1 convolutions in the internal layers.
+        kernel_valid_dims = {
+            "vertical": (rows - 1, cols),  # vertical stack
+            "horizontal": (2, cols // 2 + 1),
+        }  # horizontal stack
+
+        # Define the size of the kernel necessary to center the current pixel
+        # correctly for stride 1 convolutions in the internal layers.
+        kernel_sizes = {"vertical": (2 * rows - 3, cols), "horizontal": (3, cols)}
+
+        # Make the kernel constraint functions for stride 1 convolutions in internal
+        # layers.
+        kernel_constraints = {
+            k: _make_kernel_constraint(kernel_sizes[k], (0, v[0]), (0, v[1])) for k, v in kernel_valid_dims.items()
+        }
+
+        # Build the initial vertical stack/horizontal stack convolutional layers,
+        # as shown in Figure 1 of [2]. The receptive field of the initial vertical
+        # stack layer is a rectangular area centered above the current pixel.
+        vertical_stack_init = Conv2D(
+            kernel_size=(2 * rows - 1, cols),
+            kernel_constraint=_make_kernel_constraint((2 * rows - 1, cols), (0, rows - 1), (0, cols)),
+        )(image_input)
+
+        # In Figure 1 [2], the receptive field of the horizontal stack is
+        # illustrated as the pixels in the same row and to the left of the current
+        # pixel. [1] increases the height of this receptive field from one pixel to
+        # two (`horizontal_stack_left`) and additionally includes a subset of the
+        # row of pixels centered above the current pixel (`horizontal_stack_up`).
+        horizontal_stack_up = Conv2D(
+            kernel_size=(3, cols),
+            kernel_constraint=_make_kernel_constraint((3, cols), (0, 1), (0, cols)),
+        )(image_input)
+
+        horizontal_stack_left = Conv2D(
+            kernel_size=(3, cols),
+            kernel_constraint=_make_kernel_constraint((3, cols), (0, 2), (0, cols // 2)),
+        )(image_input)
+
+        horizontal_stack_init = keras.layers.add([horizontal_stack_up, horizontal_stack_left], dtype=dtype)
+
+        layer_stacks = {
+            "vertical": [vertical_stack_init],
+            "horizontal": [horizontal_stack_init],
+        }
+
+        # Build the downward pass of the U-net (left-hand half of Figure 2 of [1]).
+        # Each `i` iteration builds one of the highest-level blocks (identified as
+        # 'Sequence of 6 layers' in the figure, consisting of `num_resnet=5` stride-
+        # 1 layers, and one stride-2 layer that contracts the height/width
+        # dimensions). The `_` iterations build the stride 1 layers. The layers of
+        # the downward pass are stored in lists, since we'll later need them to make
+        # skip-connections to layers in the upward pass of the U-net (the skip-
+        # connections are represented by curved lines in Figure 2 [1]).
+        for i in range(self._num_hierarchies):
+            for _ in range(self._num_resnet):
+                # Build a layer shown in Figure 2 of [2]. The 'vertical' iteration
+                # builds the layers in the left half of the figure, and the 'horizontal'
+                # iteration builds the layers in the right half.
+                for stack in ["vertical", "horizontal"]:
+                    input_x = layer_stacks[stack][-1]
+                    x = activation(input_x)
+                    x = Conv2D(
+                        kernel_size=kernel_sizes[stack],
+                        kernel_constraint=kernel_constraints[stack],
+                    )(x)
+
+                    # Add the vertical-stack layer to the horizontal-stack layer
+                    if stack == "horizontal":
+                        h = activation(layer_stacks["vertical"][-1])
+                        h = Dense(self._num_filters)(h)
+                        x = keras.layers.add([h, x], dtype=dtype)
+
+                    x = activation(x)
+                    x = keras.layers.Dropout(self._dropout_p, dtype=dtype)(x)
+                    x = Conv2D(
+                        filters=2 * self._num_filters,
+                        kernel_size=kernel_sizes[stack],
+                        kernel_constraint=kernel_constraints[stack],
+                    )(x)
+
+                    if conditional_input is not None:
+                        h_projection = _build_and_apply_h_projection(conditional_input, self._num_filters, dtype=dtype)
+                        x = keras.layers.add([x, h_projection], dtype=dtype)
+
+                    x = _apply_sigmoid_gating(x)
+
+                    # Add a residual connection from the layer's input.
+                    out = keras.layers.add([input_x, x], dtype=dtype)
+                    layer_stacks[stack].append(out)
+
+            if i < self._num_hierarchies - 1:
+                # Build convolutional layers that contract the height/width dimensions
+                # on the downward pass between each set of layers (e.g. contracting from
+                # 32x32 to 16x16 in Figure 2 of [1]).
+                for stack in ["vertical", "horizontal"]:
+                    # Define kernel dimensions/masking to maintain the autoregressive property.
+                    x = layer_stacks[stack][-1]
+                    h, w = kernel_valid_dims[stack]
+                    kernel_height = 2 * h
+                    kernel_width = w + 1 if stack == "vertical" else 2 * w
+                    kernel_size = (kernel_height, kernel_width)
+                    kernel_constraint = _make_kernel_constraint(kernel_size, (0, h), (0, w))
+                    x = Conv2D(
+                        strides=(2, 2),
+                        kernel_size=kernel_size,
+                        kernel_constraint=kernel_constraint,
+                    )(x)
+                    layer_stacks[stack].append(x)
+
+        # Upward pass of the U-net (right-hand half of Figure 2 of [1]). We stored
+        # the layers of the downward pass in a list, in order to access them to make
+        # skip-connections to the upward pass. For the upward pass, we need to keep
+        # track of only the current layer, so we maintain a reference to the
+        # current layer of the horizontal/vertical stack in the `upward_pass` dict.
+        # The upward pass begins with the last layer of the downward pass.
+        upward_pass = {key: stack.pop() for key, stack in layer_stacks.items()}
+
+        # As with the downward pass, each `i` iteration builds a highest level block
+        # in Figure 2 [1], and the `_` iterations build individual layers within the
+        # block.
+        for i in range(self._num_hierarchies):
+            num_resnet = self._num_resnet if i == 0 else self._num_resnet + 1
+
+            for _ in range(num_resnet):
+                # Build a layer as shown in Figure 2 of [2], with a skip-connection
+                # from the symmetric layer in the downward pass.
+                for stack in ["vertical", "horizontal"]:
+                    input_x = upward_pass[stack]
+                    x_symmetric = layer_stacks[stack].pop()
+
+                    x = activation(input_x)
+                    x = Conv2D(
+                        kernel_size=kernel_sizes[stack],
+                        kernel_constraint=kernel_constraints[stack],
+                    )(x)
+
+                    # Include the vertical-stack layer of the upward pass in the layers
+                    # to be added to the horizontal layer.
+                    if stack == "horizontal":
+                        x_symmetric = keras.layers.Concatenate(axis=-1, dtype=dtype)(
+                            [upward_pass["vertical"], x_symmetric]
+                        )
+
+                    # Add a skip-connection from the symmetric layer in the downward
+                    # pass to the layer `x` in the upward pass.
+                    h = activation(x_symmetric)
+                    h = Dense(self._num_filters)(h)
+                    x = keras.layers.add([h, x], dtype=dtype)
+
+                    x = activation(x)
+                    x = keras.layers.Dropout(self._dropout_p, dtype=dtype)(x)
+                    x = Conv2D(
+                        filters=2 * self._num_filters,
+                        kernel_size=kernel_sizes[stack],
+                        kernel_constraint=kernel_constraints[stack],
+                    )(x)
+
+                    if conditional_input is not None:
+                        h_projection = _build_and_apply_h_projection(conditional_input, self._num_filters, dtype=dtype)
+                        x = keras.layers.add([x, h_projection], dtype=dtype)
+
+                    x = _apply_sigmoid_gating(x)
+                    upward_pass[stack] = keras.layers.add([input_x, x], dtype=dtype)
+
+            # Define deconvolutional layers that expand height/width dimensions on the
+            # upward pass (e.g. expanding from 8x8 to 16x16 in Figure 2 of [1]), with
+            # the correct kernel dimensions/masking to maintain the autoregressive
+            # property.
+            if i < self._num_hierarchies - 1:
+                for stack in ["vertical", "horizontal"]:
+                    h, w = kernel_valid_dims[stack]
+                    kernel_height = 2 * h - 2
+                    if stack == "vertical":
+                        kernel_width = w + 1
+                        kernel_constraint = _make_kernel_constraint(
+                            (kernel_height, kernel_width),
+                            (h - 2, kernel_height),
+                            (0, w),
+                        )
+                    else:
+                        kernel_width = 2 * w - 2
+                        kernel_constraint = _make_kernel_constraint(
+                            (kernel_height, kernel_width),
+                            (h - 2, kernel_height),
+                            (w - 2, kernel_width),
+                        )
+
+                    x = upward_pass[stack]
+                    x = Conv2DTranspose(
+                        kernel_size=(kernel_height, kernel_width),
+                        kernel_constraint=kernel_constraint,
+                    )(x)
+                    upward_pass[stack] = x
+
+        x_out = keras.layers.ELU(dtype=dtype)(upward_pass["horizontal"])
+
+        # Build final Dense/Reshape layers to output the correct number of
+        # parameters per pixel.
+        num_channels = tensorshape_util.as_list(image_shape)[-1]
+        num_coeffs = num_channels * (num_channels - 1) // 2  # alpha, beta, gamma in eq.3 of paper
+        num_out = num_channels * 2 + num_coeffs + 1  # mu, s + alpha, beta, gamma + 1 (mixture weight)
+        num_out_total = num_out * self._num_logistic_mix
+        params = Dense(num_out_total)(x_out)
+        params = tf.reshape(
+            params,
+            prefer_static.concat(  # [-1,H,W,nb mixtures, params per mixture]
+                [[-1], image_shape[:-1], [self._num_logistic_mix, num_out]], axis=0
+            ),
+        )
+
+        # If there is one color channel, split the parameters into a list of three
+        # output `Tensor`s: (1) component logits for the Quantized Logistic mixture
+        # distribution, (2) location parameters for each component, and (3) scale
+        # parameters for each component. If there is more than one color channel,
+        # return a fourth `Tensor` for the coefficients for the linear dependence
+        # among color channels (e.g. alpha, beta, gamma).
+        # [logits, mu, s, linear dependence]
+        splits = 3 if num_channels == 1 else [1, num_channels, num_channels, num_coeffs]
+        outputs = tf.split(params, splits, axis=-1)
+
+        # Squeeze singleton dimension from component logits
+        outputs[0] = tf.squeeze(outputs[0], axis=-1)
+
+        # Ensure scales are positive and do not collapse to near-zero
+        outputs[2] = tf.nn.softplus(outputs[2]) + tf.cast(tf.exp(-7.0), self.dtype)
+
+        inputs = image_input if conditional_input is None else [image_input, conditional_input]
+        self._network = keras.Model(inputs=inputs, outputs=outputs)
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        """Call the Pixel CNN network model.
+
+        Parameters
+        ----------
+        inputs
+            4D `Tensor` of image data with dimensions [batch size, height,
+            width, channels] or a 2-element `list`. If `list`, the first element is
+            the 4D image `Tensor` and the second element is a `Tensor` with
+            conditional input data (e.g. VAE encodings or class labels) with the
+            same leading batch dimension as the image `Tensor`.
+        training
+            `bool` or `None`. If `bool`, it controls the dropout layer,
+            where `True` implies dropout is active. If `None`, it it defaults to
+            `keras.backend.learning_phase()`
+
+        Returns
+        -------
+        outputs
+            a 3- or 4-element `list` of `Tensor`s in the following order: \
+            component_logits: 4D `Tensor` of logits for the Categorical distribution \
+            over Quantized Logistic mixture components. Dimensions are \
+            `[batch_size, height, width, num_logistic_mix]`.
+        locs
+            4D `Tensor` of location parameters for the Quantized Logistic \
+            mixture components. Dimensions are `[batch_size, height, width, \
+            num_logistic_mix, num_channels]`.
+        scales
+            4D `Tensor` of location parameters for the Quantized Logistic \
+            mixture components. Dimensions are `[batch_size, height, width, \
+            num_logistic_mix, num_channels]`.
+        coeffs
+            4D `Tensor` of coefficients for the linear dependence among \
+            color channels, included only if the image has more than one channel. \
+            Dimensions are `[batch_size, height, width, num_logistic_mix, \
+            num_coeffs]`, where `num_coeffs = num_channels * (num_channels - 1) // 2`.
+        """
+        return self._network(inputs, training=training)
+
+
+def _make_kernel_constraint(kernel_size, valid_rows, valid_columns):
+    """Make the masking function for layer kernels."""
+    mask = np.zeros(kernel_size)
+    lower, upper = valid_rows
+    left, right = valid_columns
+    mask[lower:upper, left:right] = 1.0
+    mask = mask[:, :, np.newaxis, np.newaxis]
+    return lambda x: x * mask
+
+
+def _build_and_apply_h_projection(h, num_filters, dtype):
+    """Project the conditional input."""
+    h = keras.layers.Flatten(dtype=dtype)(h)
+    h_projection = keras.layers.Dense(2 * num_filters, kernel_initializer="random_normal", dtype=dtype)(h)
+    return h_projection[..., tf.newaxis, tf.newaxis, :]
+
+
+def _apply_sigmoid_gating(x):
+    """Apply the sigmoid gating in Figure 2 of [2]."""
+    activation_tensor, gate_tensor = tf.split(x, 2, axis=-1)
+    sigmoid_gate = tf.sigmoid(gate_tensor)
+    return keras.layers.multiply([sigmoid_gate, activation_tensor], dtype=x.dtype)