docstr: 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
      `tf.keras.backend.learning_phase()`.
  return_per_feature
      `bool`. If True, return per pixel level log prob.
  
  Returns
  -------
  log_prob_values: `Tensor`.
code: 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
          `tf.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)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import numpy
  import warnings
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import logistic
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import quantized_distribution
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 3 for calling the PixelCNN obj's underlying member method code with example usage: obj._log_prob(value, conditional_input, training, return_per_feature) and returns: log_px, tf