class_docstr: Densely-connected layer class with Flipout estimator.
  
  This layer implements the Bayesian variational inference analogue to
  a dense layer by assuming the `kernel` and/or the `bias` are drawn
  from distributions. By default, the layer implements a stochastic
  forward pass via sampling from the kernel and bias posteriors,
  
  ```none
  kernel, bias ~ posterior
  outputs = tf.linalg.matmul(inputs, kernel) + bias
  ```
  
  It uses the Flipout estimator [(Wen et al., 2018)][1], which performs a Monte
  Carlo approximation of the distribution integrating over the `kernel` and
  `bias`. Flipout uses roughly twice as many floating point operations as the
  reparameterization estimator but has the advantage of significantly lower
  variance.
  
  The arguments permit separate specification of the surrogate posterior
  (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias`
  distributions.
  
  Upon being built, this layer adds losses (accessible via the `losses`
  property) representing the divergences of `kernel` and/or `bias` surrogate
  posteriors and their respective priors. When doing minibatch stochastic
  optimization, make sure to scale this loss such that it is applied just once
  per epoch (e.g. if `kl` is the sum of `losses` for each element of the batch,
  you should pass `kl / num_examples_per_epoch` to your optimizer).
  
  #### Examples
  
  We illustrate a Bayesian neural network with [variational inference](
  https://en.wikipedia.org/wiki/Variational_Bayesian_methods),
  assuming a dataset of images and length-10 one-hot `targets`.
  
  ```python
  # Using the following substitution, see:
  tfn = tfp.experimental.nn
  help(tfn.AffineVariationalReparameterization)
  BayesAffine = tfn.AffineVariationalFlipout
  ```
  
  This example uses reparameterization gradients to minimize the
  Kullback-Leibler divergence up to a constant, also known as the negative
  Evidence Lower Bound. It consists of the sum of two terms: the expected
  negative log-likelihood, which we approximate via Monte Carlo; and the KL
  divergence, which is added via regularizer terms which are arguments to the
  layer.
  
  #### References
  
  [1]: Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse. Flipout:
       Efficient Pseudo-Independent Weight Perturbations on Mini-Batches. In
       _International Conference on Learning Representations_, 2018.
       https://arxiv.org/abs/1803.04386
class_docstr_tok: ["Densely-connected","layer","class","with","Flipout","estimator",".","This","layer","implements","the","Bayesian","variational","inference","analogue","to","a","dense","layer","by","assuming","the","`","kernel","`","and\/or","the","`","bias","`","are","drawn","from","distributions",".","By","default",",","the","layer","implements","a","stochastic","forward","pass","via","sampling","from","the","kernel","and","bias","posteriors",",","``","`","none","kernel",",","bias","~","posterior","outputs","=","tf.linalg.matmul","(","inputs",",","kernel",")","+","bias","``","`","It","uses","the","Flipout","estimator","[","(","Wen","et","al.",",","2018",")","]","[","1","]",",","which","performs","a","Monte","Carlo","approximation","of","the","distribution","integrating","over","the","`","kernel","`","and","`","bias","`",".","Flipout","uses","roughly","twice","as","many","floating","point","operations","as","the","reparameterization","estimator","but","has","the","advantage","of","significantly","lower","variance",".","The","arguments","permit","separate","specification","of","the","surrogate","posterior","(","`","q","(","W|x",")","`",")",",","prior","(","`","p","(","W",")","`",")",",","and","divergence","for","both","the","`","kernel","`","and","`","bias","`","distributions",".","Upon","being","built",",","this","layer","adds","losses","(","accessible","via","the","`","losses","`","property",")","representing","the","divergences","of","`","kernel","`","and\/or","`","bias","`","surrogate","posteriors","and","their","respective","priors",".","When","doing","minibatch","stochastic","optimization",",","make","sure","to","scale","this","loss","such","that","it","is","applied","just","once","per","epoch","(","e.g",".","if","`","kl","`","is","the","sum","of","`","losses","`","for","each","element","of","the","batch",",","you","should","pass","`","kl","\/","num_examples_per_epoch","`","to","your","optimizer",")",".","#","#","#","#","Examples","We","illustrate","a","Bayesian","neural","network","with","[","variational","inference","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Variational_Bayesian_methods",")",",","assuming","a","dataset","of","images","and","length-10","one-hot","`","targets","`",".","``","`","python","#","Using","the","following","substitution",",","see",":","tfn","=","tfp.experimental.nn","help","(","tfn.AffineVariationalReparameterization",")","BayesAffine","=","tfn.AffineVariationalFlipout","``","`","This","example","uses","reparameterization","gradients","to","minimize","the","Kullback-Leibler","divergence","up","to","a","constant",",","also","known","as","the","negative","Evidence","Lower","Bound",".","It","consists","of","the","sum","of","two","terms",":","the","expected","negative","log-likelihood",",","which","we","approximate","via","Monte","Carlo",";","and","the","KL","divergence",",","which","is","added","via","regularizer","terms","which","are","arguments","to","the","layer",".","#","#","#","#","References","[","1","]",":","Yeming","Wen",",","Paul","Vicol",",","Jimmy","Ba",",","Dustin","Tran",",","and","Roger","Grosse",".","Flipout",":","Efficient","Pseudo-Independent","Weight","Perturbations","on","Mini-Batches",".","In","_International","Conference","on","Learning","Representations_",",","2018.","https",":","\/\/arxiv.org\/abs\/1803.04386"]
docstr: Constructs layer.
  
  Args:
    input_size: ...
    output_size: ...
    kernel_initializer: ...
      Default value: `None` (i.e.,
      `tfp.experimental.nn.initializers.glorot_uniform()`).
    bias_initializer: ...
      Default value: `None` (i.e., `tf_keras.initializers.zeros()`).
    make_posterior_fn: ...
      Default value:
        `tfp.experimental.nn.util.make_kernel_bias_posterior_mvn_diag`.
    make_prior_fn: ...
      Default value:
        `tfp.experimental.nn.util.make_kernel_bias_prior_spike_and_slab`.
    posterior_value_fn: ...
      Default valye: `tfd.Distribution.sample`
    unpack_weights_fn:
      Default value: `unpack_kernel_and_bias`
    dtype: ...
      Default value: `tf.float32`.
    activation_fn: ...
      Default value: `None`.
   seed: ...
      Default value: `None` (i.e., no seed).
   validate_args: ...
      Default value: `False`.
    name: ...
      Default value: `None` (i.e.,
      `'AffineVariationalFlipout'`).
docstr_tok: ["Constructs","layer",".","Args",":","input_size",":","...","output_size",":","...","kernel_initializer",":","...","Default","value",":","`","None","`","(","i.e.",",","`","tfp.experimental.nn.initializers.glorot_uniform","(",")","`",")",".","bias_initializer",":","...","Default","value",":","`","None","`","(","i.e.",",","`","tf_keras.initializers.zeros","(",")","`",")",".","make_posterior_fn",":","...","Default","value",":","`","tfp.experimental.nn.util.make_kernel_bias_posterior_mvn_diag","`",".","make_prior_fn",":","...","Default","value",":","`","tfp.experimental.nn.util.make_kernel_bias_prior_spike_and_slab","`",".","posterior_value_fn",":","...","Default","valye",":","`","tfd.Distribution.sample","`","unpack_weights_fn",":","Default","value",":","`","unpack_kernel_and_bias","`","dtype",":","...","Default","value",":","`","tf.float32","`",".","activation_fn",":","...","Default","value",":","`","None","`",".","seed",":","...","Default","value",":","`","None","`","(","i.e.",",","no","seed",")",".","validate_args",":","...","Default","value",":","`","False","`",".","name",":","...","Default","value",":","`","None","`","(","i.e.",",","`","'AffineVariationalFlipout","'","`",")","."]
code: def __init__(
      self,
      input_size,
      output_size,
      # Weights
      kernel_initializer=None,  # tfp.nn.initializers.glorot_uniform()
      bias_initializer=None,  # tf_keras.initializers.zeros()
      make_posterior_fn=kernel_bias_lib.make_kernel_bias_posterior_mvn_diag,
      make_prior_fn=kernel_bias_lib.make_kernel_bias_prior_spike_and_slab,
      posterior_value_fn=tfd.Distribution.sample,
      unpack_weights_fn=unpack_kernel_and_bias,
      dtype=tf.float32,
      # Misc
      activation_fn=None,
      seed=None,
      validate_args=False,
      name=None,
  ):
      """Constructs layer.
  
      Args:
        input_size: ...
        output_size: ...
        kernel_initializer: ...
          Default value: `None` (i.e.,
          `tfp.experimental.nn.initializers.glorot_uniform()`).
        bias_initializer: ...
          Default value: `None` (i.e., `tf_keras.initializers.zeros()`).
        make_posterior_fn: ...
          Default value:
            `tfp.experimental.nn.util.make_kernel_bias_posterior_mvn_diag`.
        make_prior_fn: ...
          Default value:
            `tfp.experimental.nn.util.make_kernel_bias_prior_spike_and_slab`.
        posterior_value_fn: ...
          Default valye: `tfd.Distribution.sample`
        unpack_weights_fn:
          Default value: `unpack_kernel_and_bias`
        dtype: ...
          Default value: `tf.float32`.
        activation_fn: ...
          Default value: `None`.
       seed: ...
          Default value: `None` (i.e., no seed).
       validate_args: ...
          Default value: `False`.
        name: ...
          Default value: `None` (i.e.,
          `'AffineVariationalFlipout'`).
      """
      self._make_posterior_fn = (
          make_posterior_fn  # For variable tracking.
      )
      self._make_prior_fn = make_prior_fn  # For variable tracking.
      batch_ndims = 0
      super(AffineVariationalFlipout, self).__init__(
          posterior=make_posterior_fn(
              [input_size, output_size],
              [output_size],
              kernel_initializer,
              bias_initializer,
              batch_ndims,
              batch_ndims,
              dtype,
          ),
          prior=make_prior_fn(
              [input_size, output_size],
              [output_size],
              kernel_initializer,
              bias_initializer,
              batch_ndims,
              batch_ndims,
              dtype,
          ),
          apply_kernel_fn=tf.matmul,
          posterior_value_fn=posterior_value_fn,
          unpack_weights_fn=unpack_weights_fn,
          dtype=dtype,
          activation_fn=activation_fn,
          seed=seed,
          validate_args=validate_args,
          name=name,
      )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.nn import layers
  from tensorflow_probability.python.experimental.nn import variational_base
  from tensorflow_probability.python.experimental.nn.util import kernel_bias
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 1 to create a new AffineVariationalFlipout object from inherited base classes: vi_lib with example: obj = AffineVariationalFlipout(input_size, output_size, kernel_initializer, bias_initializer, make_posterior_fn, make_prior_fn, posterior_value_fn, unpack_weights_fn, dtype, activation_fn, seed, validate_args, name)