class_docstr: The Pixel CNN++ distribution.
  
  Pixel CNN++ [(Salimans et al., 2017)][1] models a distribution over image
  data, parameterized by a neural network. It builds on Pixel CNN and
  Conditional Pixel CNN, as originally proposed by [(van den Oord et al.,
  2016)][2, 3]. The model expresses the joint distribution over pixels as
  the product of conditional distributions:
  `p(x|h) = prod{ p(x[i] | x[0:i], h) : i=0, ..., d }`,
  in which `p(x[i] | x[0:i], h) : i=0, ..., d` is the
  probability of the `i`-th pixel conditional on the pixels that preceded it in
  raster order (color channels in RGB order, then left to right, then top to
  bottom). `h` is optional additional data on which to condition the image
  distribution, such as class labels or VAE embeddings. The Pixel CNN++
  network enforces the dependency structure among pixels by applying a mask to
  the kernels of the convolutional layers that ensures that the values for each
  pixel depend only on other pixels up and to the left (see
  `tfd.PixelCnnNetwork`).
  
  Pixel values are modeled with a mixture of quantized logistic distributions,
  which can take on a set of distinct integer values (e.g. between 0 and 255
  for an 8-bit image).
  
  Color intensity `v` of each pixel is modeled as:
  
  `v ~ sum{q[i] * quantized_logistic(loc[i], scale[i]) : i = 0, ..., k }`,
  
  in which `k` is the number of mixture components and the `q[i]` are the
  Categorical probabilities over the components.
  
  #### Sampling
  
  Pixels are sampled one at a time, in raster order. This enforces the
  autoregressive dependency structure, in which the sample of pixel `i` is
  conditioned on the samples of pixels `1, ..., i-1`. A single color image is
  sampled as follows:
  
  ```python
  samples = random_uniform([image_height, image_width, image_channels])
  for i in image_height:
    for j in image_width:
      component_logits, locs, scales, coeffs = pixel_cnn_network(samples)
      components = Categorical(component_logits).sample()
      locs = gather(locs, components)
      scales = gather(scales, components)
  
      coef_count = 0
      channel_samples = []
      for k in image_channels:
        loc = locs[k]
        for m in range(k):
          loc += channel_samples[m] * coeffs[coef_count]
          coef_count += 1
        channel_samp = Logistic(loc, scales[k]).sample()
        channel_samples.append(channel_samp)
      samples[i, j, :] = tf.stack(channel_samples, axis=-1)
  samples = round(samples)
  ```
  
  #### Examples
  
  ```python
  
  # Build a small Pixel CNN++ model to train on MNIST.
  
  import tensorflow as tf
  import tensorflow_datasets as tfds
  import tensorflow_probability as tfp
  
  tfd = tfp.distributions
  tfk = tf_keras
  tfkl = tf_keras.layers
  
  # Load MNIST from tensorflow_datasets
  data = tfds.load('mnist')
  train_data, test_data = data['train'], data['test']
  
  def image_preprocess(x):
    x['image'] = tf.cast(x['image'], tf.float32)
    return (x['image'],)  # (input, output) of the model
  
  batch_size = 16
  train_it = train_data.map(image_preprocess).batch(batch_size).shuffle(1000)
  
  image_shape = (28, 28, 1)
  # Define a Pixel CNN network
  dist = tfd.PixelCNN(
      image_shape=image_shape,
      num_resnet=1,
      num_hierarchies=2,
      num_filters=32,
      num_logistic_mix=5,
      dropout_p=.3,
  )
  
  # Define the model input
  image_input = tfkl.Input(shape=image_shape)
  
  # Define the log likelihood for the loss fn
  log_prob = dist.log_prob(image_input)
  
  # Define the model
  model = tfk.Model(inputs=image_input, outputs=log_prob)
  model.add_loss(-tf.reduce_mean(log_prob))
  
  # Compile and train the model
  model.compile(
      optimizer=tfk.optimizers.Adam(.001),
      metrics=[])
  
  model.fit(train_it, epochs=10, verbose=True)
  
  # sample five images from the trained model
  samples = dist.sample(5)
  
  ```
  
  To train a class-conditional model:
  
  ```python
  
  data = tfds.load('mnist')
  train_data, test_data = data['train'], data['test']
  
  def image_preprocess(x):
    x['image'] = tf.cast(x['image'], tf.float32)
    # return model (inputs, outputs): inputs are (image, label) and there are no
    # outputs
    return ((x['image'], x['label']),)
  
  batch_size = 16
  train_ds = train_data.map(image_preprocess).batch(batch_size).shuffle(1000)
  optimizer = tfk.optimizers.Adam()
  
  image_shape = (28, 28, 1)
  label_shape = ()
  dist = tfd.PixelCNN(
      image_shape=image_shape,
      conditional_shape=label_shape,
      num_resnet=1,
      num_hierarchies=2,
      num_filters=32,
      num_logistic_mix=5,
      dropout_p=.3,
  )
  
  image_input = tfkl.Input(shape=image_shape)
  label_input = tfkl.Input(shape=label_shape)
  
  log_prob = dist.log_prob(image_input, conditional_input=label_input)
  
  class_cond_model = tfk.Model(
      inputs=[image_input, label_input], outputs=log_prob)
  class_cond_model.add_loss(-tf.reduce_mean(log_prob))
  class_cond_model.compile(
      optimizer=tfk.optimizers.Adam(),
      metrics=[])
  class_cond_model.fit(train_ds, epochs=10)
  
  # Take 10 samples of the digit '5'
  samples = dist.sample(10, conditional_input=5.)
  
  # Take 4 samples each of the digits '1', '2', '3'.
  # Note that when a batch of conditional input is passed, the sample shape
  # (the first argument of `dist.sample`) must have its last dimension(s) equal
  # the batch shape of the conditional input (here, (3,)).
  samples = dist.sample((4, 3), conditional_input=[1., 2., 3.])
  
  ```
  
  Note: PixelCNN may also be trained using tfp.layers.DistributionLambda;
  however, as of this writing, that method is much slower and has the
  disadvantage of calling `sample()` upon construction, which causes the
  `PixelCnnNetwork` to be initialized with random data (if data-dependent
  initialization is used).
  
  #### 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 _Neural Information Processing Systems_, 2016.
       https://arxiv.org/abs/1606.05328
  
  [3]: Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel
       Recurrent Neural Networks. In _International Conference on Machine
       Learning_, 2016. https://arxiv.org/pdf/1601.06759.pdf
class_docstr_tok: ["The","Pixel","CNN++","distribution",".","Pixel","CNN++","[","(","Salimans","et","al.",",","2017",")","]","[","1","]","models","a","distribution","over","image","data",",","parameterized","by","a","neural","network",".","It","builds","on","Pixel","CNN","and","Conditional","Pixel","CNN",",","as","originally","proposed","by","[","(","van","den","Oord","et","al.",",","2016",")","]","[","2",",","3","]",".","The","model","expresses","the","joint","distribution","over","pixels","as","the","product","of","conditional","distributions",":","`","p","(","x|h",")","=","prod","{","p","(","x","[","i","]","|","x","[","0",":","i","]",",","h",")",":","i=0",",","...",",","d","}","`",",","in","which","`","p","(","x","[","i","]","|","x","[","0",":","i","]",",","h",")",":","i=0",",","...",",","d","`","is","the","probability","of","the","`","i","`","-th","pixel","conditional","on","the","pixels","that","preceded","it","in","raster","order","(","color","channels","in","RGB","order",",","then","left","to","right",",","then","top","to","bottom",")",".","`","h","`","is","optional","additional","data","on","which","to","condition","the","image","distribution",",","such","as","class","labels","or","VAE","embeddings",".","The","Pixel","CNN++","network","enforces","the","dependency","structure","among","pixels","by","applying","a","mask","to","the","kernels","of","the","convolutional","layers","that","ensures","that","the","values","for","each","pixel","depend","only","on","other","pixels","up","and","to","the","left","(","see","`","tfd.PixelCnnNetwork","`",")",".","Pixel","values","are","modeled","with","a","mixture","of","quantized","logistic","distributions",",","which","can","take","on","a","set","of","distinct","integer","values","(","e.g",".","between","0","and","255","for","an","8-bit","image",")",".","Color","intensity","`","v","`","of","each","pixel","is","modeled","as",":","`","v","~","sum","{","q","[","i","]","*","quantized_logistic","(","loc","[","i","]",",","scale","[","i","]",")",":","i","=","0",",","...",",","k","}","`",",","in","which","`","k","`","is","the","number","of","mixture","components","and","the","`","q","[","i","]","`","are","the","Categorical","probabilities","over","the","components",".","#","#","#","#","Sampling","Pixels","are","sampled","one","at","a","time",",","in","raster","order",".","This","enforces","the","autoregressive","dependency","structure",",","in","which","the","sample","of","pixel","`","i","`","is","conditioned","on","the","samples","of","pixels","`","1",",","...",",","i-1","`",".","A","single","color","image","is","sampled","as","follows",":","``","`","python","samples","=","random_uniform","(","[","image_height",",","image_width",",","image_channels","]",")","for","i","in","image_height",":","for","j","in","image_width",":","component_logits",",","locs",",","scales",",","coeffs","=","pixel_cnn_network","(","samples",")","components","=","Categorical","(","component_logits",")",".sample","(",")","locs","=","gather","(","locs",",","components",")","scales","=","gather","(","scales",",","components",")","coef_count","=","0","channel_samples","=","[","]","for","k","in","image_channels",":","loc","=","locs","[","k","]","for","m","in","range","(","k",")",":","loc","+=","channel_samples","[","m","]","*","coeffs","[","coef_count","]","coef_count","+=","1","channel_samp","=","Logistic","(","loc",",","scales","[","k","]",")",".sample","(",")","channel_samples.append","(","channel_samp",")","samples","[","i",",","j",",",":","]","=","tf.stack","(","channel_samples",",","axis=-1",")","samples","=","round","(","samples",")","``","`","#","#","#","#","Examples","``","`","python","#","Build","a","small","Pixel","CNN++","model","to","train","on","MNIST",".","import","tensorflow","as","tf","import","tensorflow_datasets","as","tfds","import","tensorflow_probability","as","tfp","tfd","=","tfp.distributions","tfk","=","tf_keras","tfkl","=","tf_keras.layers","#","Load","MNIST","from","tensorflow_datasets","data","=","tfds.load","(","'mnist","'",")","train_data",",","test_data","=","data","[","'train","'","]",",","data","[","'test","'","]","def","image_preprocess","(","x",")",":","x","[","'image","'","]","=","tf.cast","(","x","[","'image","'","]",",","tf.float32",")","return","(","x","[","'image","'","]",",",")","#","(","input",",","output",")","of","the","model","batch_size","=","16","train_it","=","train_data.map","(","image_preprocess",")",".batch","(","batch_size",")",".shuffle","(","1000",")","image_shape","=","(","28",",","28",",","1",")","#","Define","a","Pixel","CNN","network","dist","=","tfd.PixelCNN","(","image_shape=image_shape",",","num_resnet=1",",","num_hierarchies=2",",","num_filters=32",",","num_logistic_mix=5",",","dropout_p=.3",",",")","#","Define","the","model","input","image_input","=","tfkl.Input","(","shape=image_shape",")","#","Define","the","log","likelihood","for","the","loss","fn","log_prob","=","dist.log_prob","(","image_input",")","#","Define","the","model","model","=","tfk.Model","(","inputs=image_input",",","outputs=log_prob",")","model.add_loss","(","-tf.reduce_mean","(","log_prob",")",")","#","Compile","and","train","the","model","model.compile","(","optimizer=tfk.optimizers.Adam","(",".001",")",",","metrics=","[","]",")","model.fit","(","train_it",",","epochs=10",",","verbose=True",")","#","sample","five","images","from","the","trained","model","samples","=","dist.sample","(","5",")","``","`","To","train","a","class-conditional","model",":","``","`","python","data","=","tfds.load","(","'mnist","'",")","train_data",",","test_data","=","data","[","'train","'","]",",","data","[","'test","'","]","def","image_preprocess","(","x",")",":","x","[","'image","'","]","=","tf.cast","(","x","[","'image","'","]",",","tf.float32",")","#","return","model","(","inputs",",","outputs",")",":","inputs","are","(","image",",","label",")","and","there","are","no","#","outputs","return","(","(","x","[","'image","'","]",",","x","[","'label","'","]",")",",",")","batch_size","=","16","train_ds","=","train_data.map","(","image_preprocess",")",".batch","(","batch_size",")",".shuffle","(","1000",")","optimizer","=","tfk.optimizers.Adam","(",")","image_shape","=","(","28",",","28",",","1",")","label_shape","=","(",")","dist","=","tfd.PixelCNN","(","image_shape=image_shape",",","conditional_shape=label_shape",",","num_resnet=1",",","num_hierarchies=2",",","num_filters=32",",","num_logistic_mix=5",",","dropout_p=.3",",",")","image_input","=","tfkl.Input","(","shape=image_shape",")","label_input","=","tfkl.Input","(","shape=label_shape",")","log_prob","=","dist.log_prob","(","image_input",",","conditional_input=label_input",")","class_cond_model","=","tfk.Model","(","inputs=","[","image_input",",","label_input","]",",","outputs=log_prob",")","class_cond_model.add_loss","(","-tf.reduce_mean","(","log_prob",")",")","class_cond_model.compile","(","optimizer=tfk.optimizers.Adam","(",")",",","metrics=","[","]",")","class_cond_model.fit","(","train_ds",",","epochs=10",")","#","Take","10","samples","of","the","digit","'","5'","samples","=","dist.sample","(","10",",","conditional_input=5",".",")","#","Take","4","samples","each","of","the","digits","'","1","'",",","'","2","'",",","'","3","'",".","#","Note","that","when","a","batch","of","conditional","input","is","passed",",","the","sample","shape","#","(","the","first","argument","of","`","dist.sample","`",")","must","have","its","last","dimension","(","s",")","equal","#","the","batch","shape","of","the","conditional","input","(","here",",","(","3",",",")",")",".","samples","=","dist.sample","(","(","4",",","3",")",",","conditional_input=","[","1.",",","2.",",","3",".","]",")","``","`","Note",":","PixelCNN","may","also","be","trained","using","tfp.layers.DistributionLambda",";","however",",","as","of","this","writing",",","that","method","is","much","slower","and","has","the","disadvantage","of","calling","`","sample","(",")","`","upon","construction",",","which","causes","the","`","PixelCnnNetwork","`","to","be","initialized","with","random","data","(","if","data-dependent","initialization","is","used",")",".","#","#","#","#","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","_Neural","Information","Processing","Systems_",",","2016.","https",":","\/\/arxiv.org\/abs\/1606.05328","[","3","]",":","Aaron","van","den","Oord",",","Nal","Kalchbrenner",",","and","Koray","Kavukcuoglu",".","Pixel","Recurrent","Neural","Networks",".","In","_International","Conference","on","Machine","Learning_",",","2016.","https",":","\/\/arxiv.org\/pdf\/1601.06759.pdf"]
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).
  
  Args:
    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()`.
  Returns:
    log_prob_values: `Tensor`.
docstr_tok: ["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",")",".","Args",":","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","(",")","`",".","Returns",":","log_prob_values",":","`","Tensor","`","."]
code: def _log_prob(self, value, conditional_input=None, training=None):
      """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).
  
      Args:
        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()`.
      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)
          # TODO(emilyaf) Investigate using fill_triangular/matrix multiplication
          # 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(
              transformed_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 tf.reshape(
          dist.log_prob(value), image_batch_and_conditional_shape
      )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import shift
  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 prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.layers import weight_norm
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) and returns: tf