def _resnet_block(x, filters, kernel, stride, kernel_posterior_fn):
  """Network block for ResNet."""
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  if stride != 1 or filters != x.shape[1]:
    shortcut = _projection_shortcut(x, filters, stride, kernel_posterior_fn)
  else:
    shortcut = x

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=stride,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.add([x, shortcut])
  return x

def _resnet_block ( x , filters , kernel , stride , kernel_posterior_fn ) : x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) if stride != 1 or filters != x . shape [ 1 ] : shortcut = _projection_shortcut ( x , filters , stride , kernel_posterior_fn ) else : shortcut = x x = tfp . layers . Convolution2DFlipout ( filters , kernel , strides = stride , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( x ) x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) x = tfp . layers . Convolution2DFlipout ( filters , kernel , strides = 1 , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( x ) x = tf . keras . layers . add ( [ x , shortcut ] ) return x


def make_encoder(activation, num_topics, layer_sizes):
  """Create the encoder function.

  Args:
    activation: Activation function to use.
    num_topics: The number of topics.
    layer_sizes: The number of hidden units per layer in the encoder.

  Returns:
    encoder: A `callable` mapping a bag-of-words `Tensor` to a
      `tfd.Distribution` instance over topics.
  """
  encoder_net = tf.keras.Sequential()
  for num_hidden_units in layer_sizes:
    encoder_net.add(
        tf.keras.layers.Dense(
            num_hidden_units,
            activation=activation,
            kernel_initializer=tf.compat.v1.glorot_normal_initializer()))
  encoder_net.add(
      tf.keras.layers.Dense(
          num_topics,
          activation=tf.nn.softplus,
          kernel_initializer=tf.compat.v1.glorot_normal_initializer()))

  def encoder(bag_of_words):
    net = _clip_dirichlet_parameters(encoder_net(bag_of_words))
    return tfd.Dirichlet(concentration=net,
                         name="topics_posterior")

  return encoder

def make_encoder ( activation , num_topics , layer_sizes ) : encoder_net = tf . keras . Sequential ( ) for num_hidden_units in layer_sizes : encoder_net . add ( tf . keras . layers . Dense ( num_hidden_units , activation = activation , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) encoder_net . add ( tf . keras . layers . Dense ( num_topics , activation = tf . nn . softplus , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) def encoder ( bag_of_words ) : net = _clip_dirichlet_parameters ( encoder_net ( bag_of_words ) ) return tfd . Dirichlet ( concentration = net , name = "topics_posterior" ) return encoder


def make_decoder(num_topics, num_words):
  """Create the decoder function.

  Args:
    num_topics: The number of topics.
    num_words: The number of words.

  Returns:
    decoder: A `callable` mapping a `Tensor` of encodings to a
      `tfd.Distribution` instance over words.
  """
  topics_words_logits = tf.compat.v1.get_variable(
      "topics_words_logits",
      shape=[num_topics, num_words],
      initializer=tf.compat.v1.glorot_normal_initializer())
  topics_words = tf.nn.softmax(topics_words_logits, axis=-1)

  def decoder(topics):
    word_probs = tf.matmul(topics, topics_words)
    # The observations are bag of words and therefore not one-hot. However,
    # log_prob of OneHotCategorical computes the probability correctly in
    # this case.
    return tfd.OneHotCategorical(probs=word_probs,
                                 name="bag_of_words")

  return decoder, topics_words

def make_decoder ( num_topics , num_words ) : topics_words_logits = tf . compat . v1 . get_variable ( "topics_words_logits" , shape = [ num_topics , num_words ] , initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) topics_words = tf . nn . softmax ( topics_words_logits , axis = - 1 ) def decoder ( topics ) : word_probs = tf . matmul ( topics , topics_words ) return tfd . OneHotCategorical ( probs = word_probs , name = "bag_of_words" ) return decoder , topics_words


def make_prior(num_topics, initial_value):
  """Create the prior distribution.

  Args:
    num_topics: Number of topics.
    initial_value: The starting value for the prior parameters.

  Returns:
    prior: A `callable` that returns a `tf.distribution.Distribution`
        instance, the prior distribution.
    prior_variables: A `list` of `Variable` objects, the trainable parameters
        of the prior.
  """
  def _softplus_inverse(x):
    return np.log(np.expm1(x))

  logit_concentration = tf.compat.v1.get_variable(
      "logit_concentration",
      shape=[1, num_topics],
      initializer=tf.compat.v1.initializers.constant(
          _softplus_inverse(initial_value)))
  concentration = _clip_dirichlet_parameters(
      tf.nn.softplus(logit_concentration))

  def prior():
    return tfd.Dirichlet(concentration=concentration,
                         name="topics_prior")

  prior_variables = [logit_concentration]

  return prior, prior_variables

def make_prior ( num_topics , initial_value ) : def _softplus_inverse ( x ) : return np . log ( np . expm1 ( x ) ) logit_concentration = tf . compat . v1 . get_variable ( "logit_concentration" , shape = [ 1 , num_topics ] , initializer = tf . compat . v1 . initializers . constant ( _softplus_inverse ( initial_value ) ) ) concentration = _clip_dirichlet_parameters ( tf . nn . softplus ( logit_concentration ) ) def prior ( ) : return tfd . Dirichlet ( concentration = concentration , name = "topics_prior" ) prior_variables = [ logit_concentration ] return prior , prior_variables


def sample_chain(
    num_results,
    current_state,
    previous_kernel_results=None,
    kernel=None,
    num_burnin_steps=0,
    num_steps_between_results=0,
    trace_fn=lambda current_state, kernel_results: kernel_results,
    return_final_kernel_results=False,
    parallel_iterations=10,
    name=None,
):
  """Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

  This function samples from an Markov chain at `current_state` and whose
  stationary distribution is governed by the supplied `TransitionKernel`
  instance (`kernel`).

  This function can sample from multiple chains, in parallel. (Whether or not
  there are multiple chains is dictated by the `kernel`.)

  The `current_state` can be represented as a single `Tensor` or a `list` of
  `Tensors` which collectively represent the current state.

  Since MCMC states are correlated, it is sometimes desirable to produce
  additional intermediate states, and then discard them, ending up with a set of
  states with decreased autocorrelation.  See [Owen (2017)][1]. Such "thinning"
  is made possible by setting `num_steps_between_results > 0`. The chain then
  takes `num_steps_between_results` extra steps between the steps that make it
  into the results. The extra steps are never materialized (in calls to
  `sess.run`), and thus do not increase memory requirements.

  Warning: when setting a `seed` in the `kernel`, ensure that `sample_chain`'s
  `parallel_iterations=1`, otherwise results will not be reproducible.

  In addition to returning the chain state, this function supports tracing of
  auxiliary variables used by the kernel. The traced values are selected by
  specifying `trace_fn`. By default, all kernel results are traced but in the
  future the default will be changed to no results being traced, so plan
  accordingly. See below for some examples of this feature.

  Args:
    num_results: Integer number of Markov chain draws.
    current_state: `Tensor` or Python `list` of `Tensor`s representing the
      current state(s) of the Markov chain(s).
    previous_kernel_results: A `Tensor` or a nested collection of `Tensor`s
      representing internal calculations made within the previous call to this
      function (or as returned by `bootstrap_results`).
    kernel: An instance of `tfp.mcmc.TransitionKernel` which implements one step
      of the Markov chain.
    num_burnin_steps: Integer number of chain steps to take before starting to
      collect results.
      Default value: 0 (i.e., no burn-in).
    num_steps_between_results: Integer number of chain steps between collecting
      a result. Only one out of every `num_steps_between_samples + 1` steps is
      included in the returned results.  The number of returned chain states is
      still equal to `num_results`.  Default value: 0 (i.e., no thinning).
    trace_fn: A callable that takes in the current chain state and the previous
      kernel results and return a `Tensor` or a nested collection of `Tensor`s
      that is then traced along with the chain state.
    return_final_kernel_results: If `True`, then the final kernel results are
      returned alongside the chain state and the trace specified by the
      `trace_fn`.
    parallel_iterations: The number of iterations allowed to run in parallel. It
      must be a positive integer. See `tf.while_loop` for more details.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "mcmc_sample_chain").

  Returns:
    checkpointable_states_and_trace: if `return_final_kernel_results` is
      `True`. The return value is an instance of
      `CheckpointableStatesAndTrace`.
    all_states: if `return_final_kernel_results` is `False` and `trace_fn` is
      `None`. The return value is a `Tensor` or Python list of `Tensor`s
      representing the state(s) of the Markov chain(s) at each result step. Has
      same shape as input `current_state` but with a prepended
      `num_results`-size dimension.
    states_and_trace: if `return_final_kernel_results` is `False` and
      `trace_fn` is not `None`. The return value is an instance of
      `StatesAndTrace`.

  #### Examples

  ##### Sample from a diagonal-variance Gaussian.

  I.e.,

  ```none
  for i=1..n:
    x[i] ~ MultivariateNormal(loc=0, scale=diag(true_stddev))  # likelihood
  ```

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  dims = 10
  true_stddev = np.sqrt(np.linspace(1., 3., dims))
  likelihood = tfd.MultivariateNormalDiag(loc=0., scale_diag=true_stddev)

  states = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=tf.zeros(dims),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=likelihood.log_prob,
        step_size=0.5,
        num_leapfrog_steps=2),
      trace_fn=None)

  sample_mean = tf.reduce_mean(states, axis=0)
  # ==> approx all zeros

  sample_stddev = tf.sqrt(tf.reduce_mean(
      tf.squared_difference(states, sample_mean),
      axis=0))
  # ==> approx equal true_stddev
  ```

  ##### Sampling from factor-analysis posteriors with known factors.

  I.e.,

  ```none
  # prior
  w ~ MultivariateNormal(loc=0, scale=eye(d))
  for i=1..n:
    # likelihood
    x[i] ~ Normal(loc=w^T F[i], scale=1)
  ```

  where `F` denotes factors.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Specify model.
  def make_prior(dims):
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros(dims))

  def make_likelihood(weights, factors):
    return tfd.MultivariateNormalDiag(
        loc=tf.matmul(weights, factors, adjoint_b=True))

  def joint_log_prob(num_weights, factors, x, w):
    return (make_prior(num_weights).log_prob(w) +
            make_likelihood(w, factors).log_prob(x))

  def unnormalized_log_posterior(w):
    # Posterior is proportional to: `p(W, X=x | factors)`.
    return joint_log_prob(num_weights, factors, x, w)

  # Setup data.
  num_weights = 10 # == d
  num_factors = 40 # == n
  num_chains = 100

  weights = make_prior(num_weights).sample(1)
  factors = tf.random_normal([num_factors, num_weights])
  x = make_likelihood(weights, factors).sample()

  # Sample from Hamiltonian Monte Carlo Markov Chain.

  # Get `num_results` samples from `num_chains` independent chains.
  chains_states, kernels_results = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=tf.zeros([num_chains, num_weights], name='init_weights'),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=unnormalized_log_posterior,
        step_size=0.1,
        num_leapfrog_steps=2))

  # Compute sample stats.
  sample_mean = tf.reduce_mean(chains_states, axis=[0, 1])
  # ==> approx equal to weights

  sample_var = tf.reduce_mean(
      tf.squared_difference(chains_states, sample_mean),
      axis=[0, 1])
  # ==> less than 1
  ```

  ##### Custom tracing functions.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  likelihood = tfd.Normal(loc=0., scale=1.)

  def sample_chain(trace_fn):
    return tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=0.,
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=likelihood.log_prob,
        step_size=0.5,
        num_leapfrog_steps=2),
      trace_fn=trace_fn)

  def trace_log_accept_ratio(states, previous_kernel_results):
    return previous_kernel_results.log_accept_ratio

  def trace_everything(states, previous_kernel_results):
    return previous_kernel_results

  _, log_accept_ratio = sample_chain(trace_fn=trace_log_accept_ratio)
  _, kernel_results = sample_chain(trace_fn=trace_everything)

  acceptance_prob = tf.exp(tf.minimum(log_accept_ratio_, 0.))
  # Equivalent to, but more efficient than:
  acceptance_prob = tf.exp(tf.minimum(kernel_results.log_accept_ratio_, 0.))
  ```

  #### References

  [1]: Art B. Owen. Statistically efficient thinning of a Markov chain sampler.
       _Technical Report_, 2017.
       http://statweb.stanford.edu/~owen/reports/bestthinning.pdf
  """
  if not kernel.is_calibrated:
    warnings.warn("supplied `TransitionKernel` is not calibrated. Markov "
                  "chain may not converge to intended target distribution.")
  with tf.compat.v1.name_scope(
      name, "mcmc_sample_chain",
      [num_results, num_burnin_steps, num_steps_between_results]):
    num_results = tf.convert_to_tensor(
        value=num_results, dtype=tf.int32, name="num_results")
    num_burnin_steps = tf.convert_to_tensor(
        value=num_burnin_steps, dtype=tf.int32, name="num_burnin_steps")
    num_steps_between_results = tf.convert_to_tensor(
        value=num_steps_between_results,
        dtype=tf.int32,
        name="num_steps_between_results")
    current_state = tf.nest.map_structure(
        lambda x: tf.convert_to_tensor(value=x, name="current_state"),
        current_state)
    if previous_kernel_results is None:
      previous_kernel_results = kernel.bootstrap_results(current_state)

    if trace_fn is None:
      # It simplifies the logic to use a dummy function here.
      trace_fn = lambda *args: ()
      no_trace = True
    else:
      no_trace = False
    if trace_fn is sample_chain.__defaults__[4]:
      warnings.warn("Tracing all kernel results by default is deprecated. Set "
                    "the `trace_fn` argument to None (the future default "
                    "value) or an explicit callback that traces the values "
                    "you are interested in.")

    def _trace_scan_fn(state_and_results, num_steps):
      next_state, current_kernel_results = mcmc_util.smart_for_loop(
          loop_num_iter=num_steps,
          body_fn=kernel.one_step,
          initial_loop_vars=list(state_and_results),
          parallel_iterations=parallel_iterations)
      return next_state, current_kernel_results

    (_, final_kernel_results), (all_states, trace) = mcmc_util.trace_scan(
        loop_fn=_trace_scan_fn,
        initial_state=(current_state, previous_kernel_results),
        elems=tf.one_hot(
            indices=0,
            depth=num_results,
            on_value=1 + num_burnin_steps,
            off_value=1 + num_steps_between_results,
            dtype=tf.int32),
        # pylint: disable=g-long-lambda
        trace_fn=lambda state_and_results: (state_and_results[0],
                                            trace_fn(*state_and_results)),
        # pylint: enable=g-long-lambda
        parallel_iterations=parallel_iterations)

    if return_final_kernel_results:
      return CheckpointableStatesAndTrace(
          all_states=all_states,
          trace=trace,
          final_kernel_results=final_kernel_results)
    else:
      if no_trace:
        return all_states
      else:
        return StatesAndTrace(all_states=all_states, trace=trace)

Implements Markov chain Monte Carlo via repeated TransitionKernel steps .


def deep_exponential_family(data_size, feature_size, units, shape):
  """A multi-layered topic model over a documents-by-terms matrix."""
  w2 = ed.Gamma(0.1, 0.3, sample_shape=[units[2], units[1]], name="w2")
  w1 = ed.Gamma(0.1, 0.3, sample_shape=[units[1], units[0]], name="w1")
  w0 = ed.Gamma(0.1, 0.3, sample_shape=[units[0], feature_size], name="w0")

  z2 = ed.Gamma(0.1, 0.1, sample_shape=[data_size, units[2]], name="z2")
  z1 = ed.Gamma(shape, shape / tf.matmul(z2, w2), name="z1")
  z0 = ed.Gamma(shape, shape / tf.matmul(z1, w1), name="z0")
  x = ed.Poisson(tf.matmul(z0, w0), name="x")
  return x

def deep_exponential_family ( data_size , feature_size , units , shape ) : w2 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 2 ] , units [ 1 ] ] , name = "w2" ) w1 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 1 ] , units [ 0 ] ] , name = "w1" ) w0 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 0 ] , feature_size ] , name = "w0" ) z2 = ed . Gamma ( 0.1 , 0.1 , sample_shape = [ data_size , units [ 2 ] ] , name = "z2" ) z1 = ed . Gamma ( shape , shape / tf . matmul ( z2 , w2 ) , name = "z1" ) z0 = ed . Gamma ( shape , shape / tf . matmul ( z1 , w1 ) , name = "z0" ) x = ed . Poisson ( tf . matmul ( z0 , w0 ) , name = "x" ) return x

A multi - layered topic model over a documents - by - terms matrix .


def trainable_positive_deterministic(shape, min_loc=1e-3, name=None):
  """Learnable Deterministic distribution over positive reals."""
  with tf.compat.v1.variable_scope(
      None, default_name="trainable_positive_deterministic"):
    unconstrained_loc = tf.compat.v1.get_variable("unconstrained_loc", shape)
    loc = tf.maximum(tf.nn.softplus(unconstrained_loc), min_loc)
    rv = ed.Deterministic(loc=loc, name=name)
    return rv

def trainable_positive_deterministic ( shape , min_loc = 1e-3 , name = None ) : with tf . compat . v1 . variable_scope ( None , default_name = "trainable_positive_deterministic" ) : unconstrained_loc = tf . compat . v1 . get_variable ( "unconstrained_loc" , shape ) loc = tf . maximum ( tf . nn . softplus ( unconstrained_loc ) , min_loc ) rv = ed . Deterministic ( loc = loc , name = name ) return rv

Learnable Deterministic distribution over positive reals .


def trainable_gamma(shape, min_concentration=1e-3, min_scale=1e-5, name=None):
  """Learnable Gamma via concentration and scale parameterization."""
  with tf.compat.v1.variable_scope(None, default_name="trainable_gamma"):
    unconstrained_concentration = tf.compat.v1.get_variable(
        "unconstrained_concentration",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(
            mean=0.5, stddev=0.1))
    unconstrained_scale = tf.compat.v1.get_variable(
        "unconstrained_scale",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(stddev=0.1))
    concentration = tf.maximum(tf.nn.softplus(unconstrained_concentration),
                               min_concentration)
    rate = tf.maximum(1. / tf.nn.softplus(unconstrained_scale), 1. / min_scale)
    rv = ed.Gamma(concentration=concentration, rate=rate, name=name)
    return rv

def trainable_gamma ( shape , min_concentration = 1e-3 , min_scale = 1e-5 , name = None ) : with tf . compat . v1 . variable_scope ( None , default_name = "trainable_gamma" ) : unconstrained_concentration = tf . compat . v1 . get_variable ( "unconstrained_concentration" , shape , initializer = tf . compat . v1 . initializers . random_normal ( mean = 0.5 , stddev = 0.1 ) ) unconstrained_scale = tf . compat . v1 . get_variable ( "unconstrained_scale" , shape , initializer = tf . compat . v1 . initializers . random_normal ( stddev = 0.1 ) ) concentration = tf . maximum ( tf . nn . softplus ( unconstrained_concentration ) , min_concentration ) rate = tf . maximum ( 1. / tf . nn . softplus ( unconstrained_scale ) , 1. / min_scale ) rv = ed . Gamma ( concentration = concentration , rate = rate , name = name ) return rv

Learnable Gamma via concentration and scale parameterization .


def load_nips2011_papers(path):
  """Loads NIPS 2011 conference papers.

  The NIPS 1987-2015 data set is in the form of a 11,463 x 5,812 matrix of
  per-paper word counts, containing 11,463 words and 5,811 NIPS conference
  papers (Perrone et al., 2016). We subset to papers in 2011 and words appearing
  in at least two documents and having a total word count of at least 10.

  Built from the Observations Python package.

  Args:
    path: str.
      Path to directory which either stores file or otherwise file will
      be downloaded and extracted there. Filename is `NIPS_1987-2015.csv`.

  Returns:
    bag_of_words: np.ndarray of shape [num_documents, num_words]. Each element
      denotes the number of occurrences of a specific word in a specific
      document.
    words: List of strings, denoting the words for `bag_of_words`'s columns.
  """
  path = os.path.expanduser(path)
  filename = "NIPS_1987-2015.csv"
  filepath = os.path.join(path, filename)
  if not os.path.exists(filepath):
    url = ("https://archive.ics.uci.edu/ml/machine-learning-databases/"
           "00371/NIPS_1987-2015.csv")
    if not tf.io.gfile.exists(path):
      tf.io.gfile.makedirs(path)
    print("Downloading %s to %s" % (url, filepath))
    urllib.request.urlretrieve(url, filepath)

  with open(filepath) as f:
    iterator = csv.reader(f)
    documents = next(iterator)[1:]
    words = []
    x_train = []
    for row in iterator:
      words.append(row[0])
      x_train.append(row[1:])

  x_train = np.array(x_train, dtype=np.int)

  # Subset to documents in 2011 and words appearing in at least two documents
  # and have a total word count of at least 10.
  doc_idx = [i for i, document in enumerate(documents)
             if document.startswith("2011")]
  documents = [documents[doc] for doc in doc_idx]
  x_train = x_train[:, doc_idx]
  word_idx = np.logical_and(np.sum(x_train != 0, 1) >= 2,
                            np.sum(x_train, 1) >= 10)
  words = [word for word, idx in zip(words, word_idx) if idx]
  bag_of_words = x_train[word_idx, :].T
  return bag_of_words, words


def _init_params(self, amplitude, length_scale, validate_args):
    """Shared init logic for `amplitude` and `length_scale` params.

    Args:
      amplitude: `Tensor` (or convertible) or `None` to convert, validate.
      length_scale: `Tensor` (or convertible) or `None` to convert, validate.
      validate_args: If `True`, parameters are checked for validity despite
        possibly degrading runtime performance

    Returns:
      dtype: The common `DType` of the parameters.
    """
    dtype = util.maybe_get_common_dtype(
        [amplitude, length_scale])
    if amplitude is not None:
      amplitude = tf.convert_to_tensor(
          value=amplitude, name='amplitude', dtype=dtype)
    self._amplitude = _validate_arg_if_not_none(
        amplitude, tf.compat.v1.assert_positive, validate_args)
    if length_scale is not None:
      length_scale = tf.convert_to_tensor(
          value=length_scale, name='length_scale', dtype=dtype)
    self._length_scale = _validate_arg_if_not_none(
        length_scale, tf.compat.v1.assert_positive, validate_args)
    return dtype

def _init_params ( self , amplitude , length_scale , validate_args ) : dtype = util . maybe_get_common_dtype ( [ amplitude , length_scale ] ) if amplitude is not None : amplitude = tf . convert_to_tensor ( value = amplitude , name = 'amplitude' , dtype = dtype ) self . _amplitude = _validate_arg_if_not_none ( amplitude , tf . compat . v1 . assert_positive , validate_args ) if length_scale is not None : length_scale = tf . convert_to_tensor ( value = length_scale , name = 'length_scale' , dtype = dtype ) self . _length_scale = _validate_arg_if_not_none ( length_scale , tf . compat . v1 . assert_positive , validate_args ) return dtype

Shared init logic for amplitude and length_scale params .


def _registered_kl(type_a, type_b):
  """Get the KL function registered for classes a and b."""
  hierarchy_a = tf_inspect.getmro(type_a)
  hierarchy_b = tf_inspect.getmro(type_b)
  dist_to_children = None
  kl_fn = None
  for mro_to_a, parent_a in enumerate(hierarchy_a):
    for mro_to_b, parent_b in enumerate(hierarchy_b):
      candidate_dist = mro_to_a + mro_to_b
      candidate_kl_fn = _DIVERGENCES.get((parent_a, parent_b), None)
      if not kl_fn or (candidate_kl_fn and candidate_dist < dist_to_children):
        dist_to_children = candidate_dist
        kl_fn = candidate_kl_fn
  return kl_fn

Get the KL function registered for classes a and b .


def read_image(filepath):
  """Returns an image tensor."""
  im_bytes = tf.io.read_file(filepath)
  im = tf.image.decode_image(im_bytes, channels=CHANNELS)
  im = tf.image.convert_image_dtype(im, tf.float32)
  return im

def read_image ( filepath ) : im_bytes = tf . io . read_file ( filepath ) im = tf . image . decode_image ( im_bytes , channels = CHANNELS ) im = tf . image . convert_image_dtype ( im , tf . float32 ) return im


def download_sprites():
  """Downloads the sprites data and returns the saved filepath."""
  filepath = os.path.join(FLAGS.data_dir, DATA_SPRITES_DIR)
  if not tf.io.gfile.exists(filepath):
    if not tf.io.gfile.exists(FLAGS.data_dir):
      tf.io.gfile.makedirs(FLAGS.data_dir)
    zip_name = "{}.zip".format(filepath)
    urllib.request.urlretrieve(DATA_SPRITES_URL, zip_name)
    with zipfile.ZipFile(zip_name, "r") as zip_file:
      zip_file.extractall(FLAGS.data_dir)
    tf.io.gfile.remove(zip_name)
  return filepath

def download_sprites ( ) : filepath = os . path . join ( FLAGS . data_dir , DATA_SPRITES_DIR ) if not tf . io . gfile . exists ( filepath ) : if not tf . io . gfile . exists ( FLAGS . data_dir ) : tf . io . gfile . makedirs ( FLAGS . data_dir ) zip_name = "{}.zip" . format ( filepath ) urllib . request . urlretrieve ( DATA_SPRITES_URL , zip_name ) with zipfile . ZipFile ( zip_name , "r" ) as zip_file : zip_file . extractall ( FLAGS . data_dir ) tf . io . gfile . remove ( zip_name ) return filepath

Downloads the sprites data and returns the saved filepath .


def create_character(skin, hair, top, pants):
  """Creates a character sprite from a set of attribute sprites."""
  dtype = skin.dtype
  hair_mask = tf.cast(hair[..., -1:] <= 0, dtype)
  top_mask = tf.cast(top[..., -1:] <= 0, dtype)
  pants_mask = tf.cast(pants[..., -1:] <= 0, dtype)
  char = (skin * hair_mask) + hair
  char = (char * top_mask) + top
  char = (char * pants_mask) + pants
  return char

def create_character ( skin , hair , top , pants ) : dtype = skin . dtype hair_mask = tf . cast ( hair [ ... , - 1 : ] <= 0 , dtype ) top_mask = tf . cast ( top [ ... , - 1 : ] <= 0 , dtype ) pants_mask = tf . cast ( pants [ ... , - 1 : ] <= 0 , dtype ) char = ( skin * hair_mask ) + hair char = ( char * top_mask ) + top char = ( char * pants_mask ) + pants return char

Creates a character sprite from a set of attribute sprites .


def create_seq(character, action_metadata, direction, length=8, start=0):
  """Creates a sequence.

  Args:
    character: A character sprite tensor.
    action_metadata: An action metadata tuple.
    direction: An integer representing the direction, i.e., the row
      offset within each action group corresponding to a particular
      direction.
    length: Desired length of the sequence. If this is longer than
      the number of available frames, it will roll over to the
      beginning.
    start: Index of possible frames at which to start the sequence.

  Returns:
    A sequence tensor.
  """
  sprite_start = (action_metadata[0]+direction) * FRAME_SIZE
  sprite_end = (action_metadata[0]+direction+1) * FRAME_SIZE
  sprite_line = character[sprite_start:sprite_end, ...]

  # Extract 64x64 patches that are side-by-side in the sprite, and limit
  # to the actual number of frames for the given action.
  frames = tf.stack(tf.split(sprite_line, 13, axis=1))  # 13 is a hack
  frames = frames[0:action_metadata[1]]

  # Extract a slice of the desired length.
  # NOTE: Length could be longer than the number of frames, so tile as needed.
  frames = tf.roll(frames, shift=-start, axis=0)
  frames = tf.tile(frames, [2, 1, 1, 1])  # 2 is a hack
  frames = frames[:length]
  frames = tf.cast(frames, dtype=tf.float32)
  frames.set_shape([length, FRAME_SIZE, FRAME_SIZE, CHANNELS])
  return frames

def create_seq ( character , action_metadata , direction , length = 8 , start = 0 ) : sprite_start = ( action_metadata [ 0 ] + direction ) * FRAME_SIZE sprite_end = ( action_metadata [ 0 ] + direction + 1 ) * FRAME_SIZE sprite_line = character [ sprite_start : sprite_end , ... ] frames = tf . stack ( tf . split ( sprite_line , 13 , axis = 1 ) ) frames = frames [ 0 : action_metadata [ 1 ] ] frames = tf . roll ( frames , shift = - start , axis = 0 ) frames = tf . tile ( frames , [ 2 , 1 , 1 , 1 ] ) frames = frames [ : length ] frames = tf . cast ( frames , dtype = tf . float32 ) frames . set_shape ( [ length , FRAME_SIZE , FRAME_SIZE , CHANNELS ] ) return frames


def create_random_seq(character, action_metadata, direction, length=8):
  """Creates a random sequence."""
  start = tf.random.uniform([], maxval=action_metadata[1], dtype=tf.int32)
  return create_seq(character, action_metadata, direction, length, start)

def create_random_seq ( character , action_metadata , direction , length = 8 ) : start = tf . random . uniform ( [ ] , maxval = action_metadata [ 1 ] , dtype = tf . int32 ) return create_seq ( character , action_metadata , direction , length , start )


def create_sprites_dataset(characters, actions, directions, channels=3,
                           length=8, shuffle=False, fake_data=False):
  """Creates a tf.data pipeline for the sprites dataset.

  Args:
    characters: A list of (skin, hair, top, pants) tuples containing
      relative paths to the sprite png image for each attribute.
    actions: A list of Actions.
    directions: A list of Directions.
    channels: Number of image channels to yield.
    length: Desired length of the sequences.
    shuffle: Whether or not to shuffle the characters and sequences
      start frame.
    fake_data: Boolean for whether or not to yield synthetic data.

  Returns:
    A tf.data.Dataset yielding (seq, skin label index, hair label index,
    top label index, pants label index, action label index, skin label
    name, hair label_name, top label name, pants label name, action
    label name) tuples.
  """
  if fake_data:
    dummy_image = tf.random.normal([HEIGHT, WIDTH, CHANNELS])
  else:
    basedir = download_sprites()

  action_names = [action.name for action in actions]
  action_metadata = [(action.start_row, action.frames) for action in actions]

  direction_rows = [direction.row_offset for direction in directions]

  chars = tf.data.Dataset.from_tensor_slices(characters)
  act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()
  acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()
  dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()

  if shuffle:
    chars = chars.shuffle(len(characters))

  dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))

  skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))

  def process_example(attrs, act_name, act_metadata, dir_row_offset):
    """Processes a dataset row."""
    skin_name = attrs[0]
    hair_name = attrs[1]
    top_name = attrs[2]
    pants_name = attrs[3]

    if fake_data:
      char = dummy_image
    else:
      skin = read_image(basedir + os.sep + skin_name)
      hair = read_image(basedir + os.sep + hair_name)
      top = read_image(basedir + os.sep + top_name)
      pants = read_image(basedir + os.sep + pants_name)
      char = create_character(skin, hair, top, pants)

    if shuffle:
      seq = create_random_seq(char, act_metadata, dir_row_offset, length)
    else:
      seq = create_seq(char, act_metadata, dir_row_offset, length)
    seq = seq[..., :channels]  # limit output channels

    skin_idx = skin_table.lookup(skin_name)
    hair_idx = hair_table.lookup(hair_name)
    top_idx = top_table.lookup(top_name)
    pants_idx = pants_table.lookup(pants_name)
    act_idx = action_table.lookup(act_name)

    return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,
            skin_name, hair_name, top_name, pants_name, act_name)

  dataset = dataset.map(process_example)
  return dataset

def create_sprites_dataset ( characters , actions , directions , channels = 3 , length = 8 , shuffle = False , fake_data = False ) : if fake_data : dummy_image = tf . random . normal ( [ HEIGHT , WIDTH , CHANNELS ] ) else : basedir = download_sprites ( ) action_names = [ action . name for action in actions ] action_metadata = [ ( action . start_row , action . frames ) for action in actions ] direction_rows = [ direction . row_offset for direction in directions ] chars = tf . data . Dataset . from_tensor_slices ( characters ) act_names = tf . data . Dataset . from_tensor_slices ( action_names ) . repeat ( ) acts_metadata = tf . data . Dataset . from_tensor_slices ( action_metadata ) . repeat ( ) dir_rows = tf . data . Dataset . from_tensor_slices ( direction_rows ) . repeat ( ) if shuffle : chars = chars . shuffle ( len ( characters ) ) dataset = tf . data . Dataset . zip ( ( chars , act_names , acts_metadata , dir_rows ) ) skin_table = tf . contrib . lookup . index_table_from_tensor ( sorted ( SKIN_COLORS ) ) hair_table = tf . contrib . lookup . index_table_from_tensor ( sorted ( HAIRSTYLES ) ) top_table = tf . contrib . lookup . index_table_from_tensor ( sorted ( TOPS ) ) pants_table = tf . contrib . lookup . index_table_from_tensor ( sorted ( PANTS ) ) action_table = tf . contrib . lookup . index_table_from_tensor ( sorted ( action_names ) ) def process_example ( attrs , act_name , act_metadata , dir_row_offset ) : skin_name = attrs [ 0 ] hair_name = attrs [ 1 ] top_name = attrs [ 2 ] pants_name = attrs [ 3 ] if fake_data : char = dummy_image else : skin = read_image ( basedir + os . sep + skin_name ) hair = read_image ( basedir + os . sep + hair_name ) top = read_image ( basedir + os . sep + top_name ) pants = read_image ( basedir + os . sep + pants_name ) char = create_character ( skin , hair , top , pants ) if shuffle : seq = create_random_seq ( char , act_metadata , dir_row_offset , length ) else : seq = create_seq ( char , act_metadata , dir_row_offset , length ) seq = seq [ ... , : channels ] skin_idx = skin_table . lookup ( skin_name ) hair_idx = hair_table . lookup ( hair_name ) top_idx = top_table . lookup ( top_name ) pants_idx = pants_table . lookup ( pants_name ) act_idx = action_table . lookup ( act_name ) return ( seq , skin_idx , hair_idx , top_idx , pants_idx , act_idx , skin_name , hair_name , top_name , pants_name , act_name ) dataset = dataset . map ( process_example ) return dataset

Creates a tf . data pipeline for the sprites dataset .


def _maybe_validate_distributions(distributions, dtype_override, validate_args):
  """Checks that `distributions` satisfies all assumptions."""
  assertions = []

  if not _is_iterable(distributions) or not distributions:
    raise ValueError('`distributions` must be a list of one or more '
                     'distributions.')

  if dtype_override is None:
    dts = [
        dtype_util.base_dtype(d.dtype)
        for d in distributions
        if d.dtype is not None
    ]
    if dts[1:] != dts[:-1]:
      raise TypeError('Distributions must have same dtype; found: {}.'.format(
          set(dtype_util.name(dt) for dt in dts)))

  # Validate event_ndims.
  for d in distributions:
    if tensorshape_util.rank(d.event_shape) is not None:
      if tensorshape_util.rank(d.event_shape) != 1:
        raise ValueError('`Distribution` must be vector variate, '
                         'found event nimds: {}.'.format(
                             tensorshape_util.rank(d.event_shape)))
    elif validate_args:
      assertions.append(
          assert_util.assert_equal(
              1, tf.size(input=d.event_shape_tensor()),
              message='`Distribution` must be vector variate.'))

  batch_shapes = [d.batch_shape for d in distributions]
  if all(tensorshape_util.is_fully_defined(b) for b in batch_shapes):
    if batch_shapes[1:] != batch_shapes[:-1]:
      raise ValueError('Distributions must have the same `batch_shape`; '
                       'found: {}.'.format(batch_shapes))
  elif validate_args:
    batch_shapes = [
        tensorshape_util.as_list(d.batch_shape)  # pylint: disable=g-complex-comprehension
        if tensorshape_util.is_fully_defined(d.batch_shape) else
        d.batch_shape_tensor() for d in distributions
    ]
    assertions.extend(
        assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
            b1, b2,
            message='Distribution `batch_shape`s must be identical.')
        for b1, b2 in zip(batch_shapes[1:], batch_shapes[:-1]))

  return assertions

Checks that distributions satisfies all assumptions .


def _flatten_summand_list(kernels):
  """Flatten a list of kernels which may contain _SumKernel instances.

  Args:
    kernels: Python list of `PositiveSemidefiniteKernel` instances

  Returns:
    Python list containing the elements of kernels, with any _SumKernel
    instances replaced by their `kernels` property contents.
  """
  flattened = []
  for k in kernels:
    if isinstance(k, _SumKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened

def _flatten_summand_list ( kernels ) : flattened = [ ] for k in kernels : if isinstance ( k , _SumKernel ) : flattened += k . kernels else : flattened . append ( k ) return flattened

Flatten a list of kernels which may contain _SumKernel instances .


def _flatten_multiplicand_list(kernels):
  """Flatten a list of kernels which may contain _ProductKernel instances.

  Args:
    kernels: Python list of `PositiveSemidefiniteKernel` instances

  Returns:
    Python list containing the elements of kernels, with any _ProductKernel
    instances replaced by their `kernels` property contents.
  """
  flattened = []
  for k in kernels:
    if isinstance(k, _ProductKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened

def _flatten_multiplicand_list ( kernels ) : flattened = [ ] for k in kernels : if isinstance ( k , _ProductKernel ) : flattened += k . kernels else : flattened . append ( k ) return flattened

Flatten a list of kernels which may contain _ProductKernel instances .


def build_fake_data():
  """Build fake CIFAR10-style data for unit testing."""
  num_examples = 10
  x_train = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_train = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  x_test = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_test = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  return (x_train, y_train), (x_test, y_test)

def build_fake_data ( ) : num_examples = 10 x_train = np . random . rand ( num_examples , * IMAGE_SHAPE ) . astype ( np . float32 ) y_train = np . random . permutation ( np . arange ( num_examples ) ) . astype ( np . int32 ) x_test = np . random . rand ( num_examples , * IMAGE_SHAPE ) . astype ( np . float32 ) y_test = np . random . permutation ( np . arange ( num_examples ) ) . astype ( np . int32 ) return ( x_train , y_train ) , ( x_test , y_test )

Build fake CIFAR10 - style data for unit testing .


def count_integers(arr,
                   weights=None,
                   minlength=None,
                   maxlength=None,
                   axis=None,
                   dtype=tf.int32,
                   name=None):
  """Counts the number of occurrences of each value in an integer array `arr`.

  Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
  dimensions to reduce over.  With
    `~axis = [i for i in range(arr.ndim) if i not in axis]`,
  this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.

  If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
  if `arr` is non-empty, and 0 otherwise.
  If `weights` are non-None, then index `i` of the output stores the sum of the
  value in `weights` at each index where the corresponding value in `arr` is
  `i`.

  Args:
    arr: An `int32` `Tensor` of non-negative values.
    weights: If non-None, must be the same shape as arr. For each value in
      `arr`, the bin will be incremented by the corresponding weight instead of
      1.
    minlength: If given, ensures the output has length at least `minlength`,
      padding with zeros at the end if necessary.
    maxlength: If given, skips values in `arr` that are equal or greater than
      `maxlength`, ensuring that the output has length at most `maxlength`.
    axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
      dimensions in `arr` to reduce over.
      `Default value:` `None`, meaning reduce over all dimensions.
    dtype: If `weights` is None, determines the type of the output bins.
    name: A name scope for the associated operations (optional).

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
  """
  with tf.compat.v1.name_scope(
      name, 'count_integers', values=[arr, weights, minlength, maxlength,
                                      axis]):
    if axis is None:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='arr')
    arr_ndims = _get_static_ndims(arr, expect_static=True)

    axis = _make_static_axis_non_negative_list(axis, arr_ndims)

    # ~axis from docstring.  Dims in arr that are not in axis.
    not_axis = sorted(set(range(arr_ndims)).difference(axis))

    # If we're reducing over everything, just use standard bincount.
    if not not_axis:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    # Move dims in ~axis to the left, so we can tf.map_fn bincount over them,
    # Producing counts for every index I in ~axis.
    # Thus, flat_arr is not totally flat, it just has the dims in ~axis
    # flattened.
    flat_arr = _move_dims_to_flat_end(arr, not_axis, arr_ndims, right_end=False)

    # tf.map_fn over dim 0.
    if weights is None:

      def one_bincount(arr_slice):
        return tf.math.bincount(
            arr_slice,
            weights=None,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)
    else:
      weights = tf.convert_to_tensor(value=weights, name='weights')
      _get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)
      flat_weights = _move_dims_to_flat_end(
          weights, not_axis, arr_ndims, right_end=False)

      def one_bincount(arr_and_weights_slices):
        arr_slice, weights_slice = arr_and_weights_slices
        return tf.math.bincount(
            arr_slice,
            weights=weights_slice,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(
          one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)

    # flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0.
    # bincount needs to have the K bins in axis 0, so transpose...
    flat_counts_t = tf.transpose(a=flat_counts, perm=[1, 0])

    # Throw in this assert, to ensure shape assumptions are correct.
    _get_static_ndims(flat_counts_t, expect_ndims=2, expect_static=True)

    # not_axis_shape = arr.shape[~axis]
    not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    out_shape = tf.concat([[-1], not_axis_shape], axis=0)

    return tf.reshape(flat_counts_t, out_shape)

def count_integers ( arr , weights = None , minlength = None , maxlength = None , axis = None , dtype = tf . int32 , name = None ) : with tf . compat . v1 . name_scope ( name , 'count_integers' , values = [ arr , weights , minlength , maxlength , axis ] ) : if axis is None : return tf . math . bincount ( arr , weights = weights , minlength = minlength , maxlength = maxlength , dtype = dtype ) arr = tf . convert_to_tensor ( value = arr , dtype = tf . int32 , name = 'arr' ) arr_ndims = _get_static_ndims ( arr , expect_static = True ) axis = _make_static_axis_non_negative_list ( axis , arr_ndims ) not_axis = sorted ( set ( range ( arr_ndims ) ) . difference ( axis ) ) if not not_axis : return tf . math . bincount ( arr , weights = weights , minlength = minlength , maxlength = maxlength , dtype = dtype ) flat_arr = _move_dims_to_flat_end ( arr , not_axis , arr_ndims , right_end = False ) if weights is None : def one_bincount ( arr_slice ) : return tf . math . bincount ( arr_slice , weights = None , minlength = minlength , maxlength = maxlength , dtype = dtype ) flat_counts = tf . map_fn ( one_bincount , elems = flat_arr , dtype = dtype ) else : weights = tf . convert_to_tensor ( value = weights , name = 'weights' ) _get_static_ndims ( weights , expect_static = True , expect_ndims = arr_ndims ) flat_weights = _move_dims_to_flat_end ( weights , not_axis , arr_ndims , right_end = False ) def one_bincount ( arr_and_weights_slices ) : arr_slice , weights_slice = arr_and_weights_slices return tf . math . bincount ( arr_slice , weights = weights_slice , minlength = minlength , maxlength = maxlength , dtype = dtype ) flat_counts = tf . map_fn ( one_bincount , elems = [ flat_arr , flat_weights ] , dtype = weights . dtype ) flat_counts_t = tf . transpose ( a = flat_counts , perm = [ 1 , 0 ] ) _get_static_ndims ( flat_counts_t , expect_ndims = 2 , expect_static = True ) not_axis_shape = tf . gather ( tf . shape ( input = arr ) , indices = not_axis ) out_shape = tf . concat ( [ [ - 1 ] , not_axis_shape ] , axis = 0 ) return tf . reshape ( flat_counts_t , out_shape )

Counts the number of occurrences of each value in an integer array arr .


def find_bins(x,
              edges,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Bin values into discrete intervals.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function returns `bins`, such that:
  `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have
      `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
      designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
      corresponding dimensions of `x`.
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      This effects the output values when `x` is below/above the intervals,
      which will be `-1/K+1` for `int` types and `NaN` for `float`s.
      At indices where `x` is `NaN`, the output values will be `0` for `int`
      types and `NaN` for floats.
    name:  A Python string name to prepend to created ops. Default: 'find_bins'

  Returns:
    bins: `Tensor` with same `shape` as `x` and `dtype`.
      Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
      `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Raises:
    ValueError:  If `edges.shape[0]` is determined to be less than 2.

  #### Examples

  Cut a `1-D` array

  ```python
  x = [0., 5., 6., 10., 20.]
  edges = [0., 5., 10.]
  tfp.stats.find_bins(x, edges)
  ==> [0., 0., 1., 1., np.nan]
  ```

  Cut `x` into its deciles

  ```python
  x = tf.random_uniform(shape=(100, 200))
  decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
  bins = tfp.stats.find_bins(x, edges=decile_edges)
  bins.shape
  ==> (100, 200)
  tf.reduce_mean(bins == 0.)
  ==> approximately 0.1
  tf.reduce_mean(bins == 1.)
  ==> approximately 0.1
  ```

  """
  # TFP users may be surprised to see the "action" in the leftmost dim of
  # edges, rather than the rightmost (event) dim.  Why?
  # 1. Most likely you created edges by getting quantiles over samples, and
  #    quantile/percentile return these edges in the leftmost (sample) dim.
  # 2. Say you have event_shape = [5], then we expect the bin will be different
  #    for all 5 events, so the index of the bin should not be in the event dim.
  with tf.compat.v1.name_scope(
      name, default_name='find_bins', values=[x, edges]):
    in_type = dtype_util.common_dtype([x, edges],
                                      preferred_dtype=tf.float32)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_type)
    x = tf.convert_to_tensor(value=x, name='x', dtype=in_type)

    if (tf.compat.dimension_value(edges.shape[0]) is not None and
        tf.compat.dimension_value(edges.shape[0]) < 2):
      raise ValueError(
          'First dimension of `edges` must have length > 1 to index 1 or '
          'more bin. Found: {}'.format(edges.shape))

    flattening_x = edges.shape.ndims == 1 and x.shape.ndims > 1

    if flattening_x:
      x_orig_shape = tf.shape(input=x)
      x = tf.reshape(x, [-1])

    if dtype is None:
      dtype = in_type
    dtype = tf.as_dtype(dtype)

    # Move first dims into the rightmost.
    x_permed = distribution_util.rotate_transpose(x, shift=-1)
    edges_permed = distribution_util.rotate_transpose(edges, shift=-1)

    # If...
    #   x_permed = [0, 1, 6., 10]
    #   edges = [0, 5, 10.]
    #   ==> almost_output = [0, 1, 2, 2]
    searchsorted_type = dtype if dtype in [tf.int32, tf.int64] else None
    almost_output_permed = tf.searchsorted(
        sorted_sequence=edges_permed,
        values=x_permed,
        side='right',
        out_type=searchsorted_type)
    # Move the rightmost dims back to the leftmost.
    almost_output = tf.cast(
        distribution_util.rotate_transpose(almost_output_permed, shift=1),
        dtype)

    # In above example, we want [0, 0, 1, 1], so correct this here.
    bins = tf.clip_by_value(almost_output - 1, tf.cast(0, dtype),
                            tf.cast(tf.shape(input=edges)[0] - 2, dtype))

    if not extend_lower_interval:
      low_fill = np.nan if dtype.is_floating else -1
      bins = tf.where(x < tf.expand_dims(edges[0], 0),
                      tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),
                      bins)

    if not extend_upper_interval:
      up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[0] - 1
      bins = tf.where(x > tf.expand_dims(edges[-1], 0),
                      tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)

    if flattening_x:
      bins = tf.reshape(bins, x_orig_shape)

    return bins


def histogram(x,
              edges,
              axis=None,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Count how often `x` falls in intervals defined by `edges`.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function counts how often `x` falls into each interval.

  Values of `x` outside of the intervals cause errors.  Consider using
  `extend_lower_interval`, `extend_upper_interval` to deal with this.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not
      `None`, must have statically known number of dimensions. The
      `axis` kwarg determines which dimensions index iid samples.
      Other dimensions of `x` index "events" for which we will compute different
      histograms.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have `edges.shape[1:]` the same
      as the dimensions of `x` excluding `axis`.
      If `rank(edges) > 1`, `edges[k]` designates a shape `edges.shape[1:]`
      `Tensor` of interval edges for the corresponding dimensions of `x`.
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant
      values. The axis in `x` that index iid samples.
      `Default value:` `None` (treat every dimension as sample dimension).
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
    name:  A Python string name to prepend to created ops.
      `Default value:` 'histogram'

  Returns:
    counts: `Tensor` of type `dtype` and, with
      `~axis = [i for i in range(arr.ndim) if i not in axis]`,
      `counts.shape = [edges.shape[0]] + x.shape[~axis]`.
      With `I` a multi-index into `~axis`, `counts[k][I]` is the number of times
      event(s) fell into the `kth` interval of `edges`.

  #### Examples

  ```python
  # x.shape = [1000, 2]
  # x[:, 0] ~ Uniform(0, 1), x[:, 1] ~ Uniform(1, 2).
  x = tf.stack([tf.random_uniform([1000]), 1 + tf.random_uniform([1000])],
               axis=-1)

  # edges ==> bins [0, 0.5), [0.5, 1.0), [1.0, 1.5), [1.5, 2.0].
  edges = [0., 0.5, 1.0, 1.5, 2.0]

  tfp.stats.histogram(x, edges)
  ==> approximately [500, 500, 500, 500]

  tfp.stats.histogram(x, edges, axis=0)
  ==> approximately [[500, 500, 0, 0], [0, 0, 500, 500]]
  ```

  """
  with tf.compat.v1.name_scope(name, 'histogram', values=[x, edges, axis]):

    # Tensor conversions.
    in_dtype = dtype_util.common_dtype([x, edges], preferred_dtype=tf.float32)

    x = tf.convert_to_tensor(value=x, name='x', dtype=in_dtype)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_dtype)

    # Move dims in axis to the left end as one flattened dim.
    # After this, x.shape = [n_samples] + E.
    if axis is None:
      x = tf.reshape(x, shape=[-1])
    else:
      x_ndims = _get_static_ndims(
          x, expect_static=True, expect_ndims_at_least=1)
      axis = _make_static_axis_non_negative_list(axis, x_ndims)
      if not axis:
        raise ValueError('`axis` cannot be empty.  Found: {}'.format(axis))
      x = _move_dims_to_flat_end(x, axis, x_ndims, right_end=False)

    # bins.shape = x.shape = [n_samples] + E,
    # and bins[i] is a shape E Tensor of the bins that sample `i` fell into.
    # E is the "event shape", which is [] if axis is None.
    bins = find_bins(
        x,
        edges=edges,
        # If not extending intervals, then values outside the edges will return
        # -1, which gives an error when fed to bincount.
        extend_lower_interval=extend_lower_interval,
        extend_upper_interval=extend_upper_interval,
        dtype=tf.int32)

    # TODO(b/124015136) Use standard tf.math.bincount once it supports `axis`.
    counts = count_integers(
        bins,
        # Ensure we get correct output, even if x did not fall into every bin
        minlength=tf.shape(input=edges)[0] - 1,
        maxlength=tf.shape(input=edges)[0] - 1,
        axis=0,
        dtype=dtype or in_dtype)
    n_edges = tf.compat.dimension_value(edges.shape[0])
    if n_edges is not None:
      counts.set_shape(
          tf.TensorShape([n_edges - 1]).concatenate(counts.shape[1:]))
    return counts

def
histogram
(
x
,
edges
,
axis
=
,
extend_lower_interval
=
False
,
extend_upper_interval
=
False
,
dtype
=
,
name
=
)
:
with
tf
.
compat
.
v1
.
name_scope
(
name
,
'histogram'
,
values
=
[
x
,
edges
,
axis
]
)
:
in_dtype
=
dtype_util
.
common_dtype
(
[
x
,
edges
]
,
preferred_dtype
=
tf
.
float32
)
x
=
tf
.
convert_to_tensor
(
value
=
x
,
name
=
'x'
,
dtype
=
in_dtype
)
edges
=
tf
.
convert_to_tensor
(
value
=
edges
,
name
=
'edges'
,
dtype
=
in_dtype
)
if
axis
is
:
x
=
tf
.
reshape
(
x
,
shape
=
[
1
]
)
else
:
x_ndims
=
_get_static_ndims
(
x
,
expect_static
=
True
,
expect_ndims_at_least
=
1
)
axis
=
_make_static_axis_non_negative_list
(
axis
,
x_ndims
)
if
not
axis
:
raise
ValueError
(
'`axis` cannot be empty.  Found: {}'
.
format
(
axis
)
)
x
=
_move_dims_to_flat_end
(
x
,
axis
,
x_ndims
,
right_end
=
False
)
bins
=
find_bins
(
x
,
edges
=
edges
,
extend_lower_interval
=
extend_lower_interval
,
extend_upper_interval
=
extend_upper_interval
,
dtype
=
tf
.
int32
)
counts
=
count_integers
(
bins
,
minlength
=
tf
.
shape
(
input
=
edges
)
[
0
]
1
,
maxlength
=
tf
.
shape
(
input
=
edges
)
[
0
]
1
,
axis
=
0
,
dtype
=
dtype
or
in_dtype
)
n_edges
=
tf
.
compat
.
dimension_value
(
edges
.
shape
[
0
]
)
if
n_edges
is
not
:
counts
.
set_shape
(
tf
.
TensorShape
(
[
n_edges
1
]
)
.
concatenate
(
counts
.
shape
[
1
:
]
)
)
return
counts

Count how often x falls in intervals defined by edges .


def quantiles(x,
              num_quantiles,
              axis=None,
              interpolation=None,
              keep_dims=False,
              validate_args=False,
              name=None):
  """Compute quantiles of `x` along `axis`.

  The quantiles of a distribution are cut points dividing the range into
  intervals with equal probabilities.

  Given a vector `x` of samples, this function estimates the cut points by
  returning `num_quantiles + 1` cut points, `(c0, ..., cn)`, such that, roughly
  speaking, equal number of sample points lie in the `num_quantiles` intervals
  `[c0, c1), [c1, c2), ..., [c_{n-1}, cn]`.  That is,

  * About `1 / n` fraction of the data lies in `[c_{k-1}, c_k)`, `k = 1, ..., n`
  * About `k / n` fraction of the data lies below `c_k`.
  * `c0` is the sample minimum and `cn` is the maximum.

  The exact number of data points in each interval depends on the size of
  `x` (e.g. whether the size is divisible by `n`) and the `interpolation` kwarg.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
      `x` must have statically known number of dimensions.
    num_quantiles:  Scalar `integer` `Tensor`.  The number of intervals the
      returned `num_quantiles + 1` cut points divide the range into.
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
      axis that index independent samples over which to return the desired
      percentile.  If `None` (the default), treat every dimension as a sample
      dimension, returning a scalar.
    interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
      Default value: 'nearest'.  This specifies the interpolation method to
      use when the fractions `k / n` lie between two data points `i < j`:
        * linear: i + (j - i) * fraction, where fraction is the fractional part
          of the index surrounded by i and j.
        * lower: `i`.
        * higher: `j`.
        * nearest: `i` or `j`, whichever is nearest.
        * midpoint: (i + j) / 2. `linear` and `midpoint` interpolation do not
          work with integer dtypes.
    keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
      If `False`, the last dimension is removed from the output shape.
    validate_args:  Whether to add runtime checks of argument validity. If
      False, and arguments are incorrect, correct behavior is not guaranteed.
    name:  A Python string name to give this `Op`.  Default is 'percentile'

  Returns:
    cut_points:  A `rank(x) + 1 - len(axis)` dimensional `Tensor` with same
    `dtype` as `x` and shape `[num_quantiles + 1, ...]` where the trailing shape
    is that of `x` without the dimensions in `axis` (unless `keep_dims is True`)

  Raises:
    ValueError:  If argument 'interpolation' is not an allowed type.
    ValueError:  If interpolation type not compatible with `dtype`.

  #### Examples

  ```python
  # Get quartiles of x with various interpolation choices.
  x = [0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='nearest')
  ==> [  0.,   2.,   5.,   8.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='linear')
  ==> [  0. ,   2.5,   5. ,   7.5,  10. ]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='lower')
  ==> [  0.,   2.,   5.,   7.,  10.]

  # Get deciles of columns of an R x C data set.
  data = load_my_columnar_data(...)
  tfp.stats.quantiles(data, num_quantiles=10)
  ==> Shape [11, C] Tensor
  ```

  """
  with tf.compat.v1.name_scope(
      name, 'quantiles', values=[x, num_quantiles, axis]):
    x = tf.convert_to_tensor(value=x, name='x')
    return percentile(
        x,
        q=tf.linspace(
            # percentile casts q to float64 before using it...so may as well use
            # float64 here. Note that  using x.dtype won't work with linspace
            # if x is integral type (which is anothe motivation for hard-coding
            # float64).
            tf.convert_to_tensor(value=0, dtype=tf.float64),
            tf.convert_to_tensor(value=100, dtype=tf.float64),
            num=num_quantiles + 1),
        axis=axis,
        interpolation=interpolation,
        keep_dims=keep_dims,
        validate_args=validate_args,
        preserve_gradients=False)

def quantiles ( x , num_quantiles , axis = None , interpolation = None , keep_dims = False , validate_args = False , name = None ) : with tf . compat . v1 . name_scope ( name , 'quantiles' , values = [ x , num_quantiles , axis ] ) : x = tf . convert_to_tensor ( value = x , name = 'x' ) return percentile ( x , q = tf . linspace ( tf . convert_to_tensor ( value = 0 , dtype = tf . float64 ) , tf . convert_to_tensor ( value = 100 , dtype = tf . float64 ) , num = num_quantiles + 1 ) , axis = axis , interpolation = interpolation , keep_dims = keep_dims , validate_args = validate_args , preserve_gradients = False )


def _get_static_ndims(x,
                      expect_static=False,
                      expect_ndims=None,
                      expect_ndims_no_more_than=None,
                      expect_ndims_at_least=None):
  """Get static number of dimensions and assert that some expectations are met.

  This function returns the number of dimensions 'ndims' of x, as a Python int.

  The optional expect arguments are used to check the ndims of x, but this is
  only done if the static ndims of x is not None.

  Args:
    x:  A Tensor.
    expect_static:  Expect `x` to have statically defined `ndims`.
    expect_ndims:  Optional Python integer.  If provided, assert that x has
      number of dimensions equal to this.
    expect_ndims_no_more_than:  Optional Python integer.  If provided, assert
      that x has no more than this many dimensions.
    expect_ndims_at_least:  Optional Python integer.  If provided, assert that x
      has at least this many dimensions.

  Returns:
    ndims:  A Python integer.

  Raises:
    ValueError:  If any of the expectations above are violated.
  """
  ndims = x.shape.ndims
  if ndims is None:
    shape_const = tf.get_static_value(tf.shape(input=x))
    if shape_const is not None:
      ndims = shape_const.ndim

  if ndims is None:
    if expect_static:
      raise ValueError(
          'Expected argument `x` to have statically defined `ndims`.  Found: ' %
          x)
    return

  if expect_ndims is not None:
    ndims_message = ('Expected argument `x` to have ndims %s.  Found tensor %s'
                     % (expect_ndims, x))
    if ndims != expect_ndims:
      raise ValueError(ndims_message)

  if expect_ndims_at_least is not None:
    ndims_at_least_message = (
        'Expected argument `x` to have ndims >= %d.  Found tensor %s' %
        (expect_ndims_at_least, x))
    if ndims < expect_ndims_at_least:
      raise ValueError(ndims_at_least_message)

  if expect_ndims_no_more_than is not None:
    ndims_no_more_than_message = (
        'Expected argument `x` to have ndims <= %d.  Found tensor %s' %
        (expect_ndims_no_more_than, x))
    if ndims > expect_ndims_no_more_than:
      raise ValueError(ndims_no_more_than_message)

  return ndims

Get static number of dimensions and assert that some expectations are met .


def _insert_back_keep_dims(x, axis):
  """Insert the dims in `axis` back as singletons after being removed.

  Args:
    x:  `Tensor`.
    axis:  Python list of integers.

  Returns:
    `Tensor` with same values as `x`, but additional singleton dimensions.
  """
  for i in sorted(axis):
    x = tf.expand_dims(x, axis=i)
  return x

def _insert_back_keep_dims ( x , axis ) : for i in sorted ( axis ) : x = tf . expand_dims ( x , axis = i ) return x

Insert the dims in axis back as singletons after being removed .


def _make_static_axis_non_negative_list(axis, ndims):
  """Convert possibly negatively indexed axis to non-negative list of ints.

  Args:
    axis:  Integer Tensor.
    ndims:  Number of dimensions into which axis indexes.

  Returns:
    A list of non-negative Python integers.

  Raises:
    ValueError: If `axis` is not statically defined.
  """
  axis = distribution_util.make_non_negative_axis(axis, ndims)

  axis_const = tf.get_static_value(axis)
  if axis_const is None:
    raise ValueError(
        'Expected argument `axis` to be statically available.  Found: %s' %
        axis)

  # Make at least 1-D.
  axis = axis_const + np.zeros([1], dtype=axis_const.dtype)

  return list(int(dim) for dim in axis)

Convert possibly negatively indexed axis to non - negative list of ints .


def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):
  """Move dims corresponding to `axis` in `x` to the end, then flatten.

  Args:
    x: `Tensor` with shape `[B0,B1,...,Bb]`.
    axis:  Python list of indices into dimensions of `x`.
    x_ndims:  Python integer holding number of dimensions in `x`.
    right_end:  Python bool.  Whether to move dims to the right end (else left).

  Returns:
    `Tensor` with value from `x` and dims in `axis` moved to end into one single
      dimension.
  """

  if not axis:
    return x

  # Suppose x.shape = [a, b, c, d]
  # Suppose axis = [1, 3]

  # other_dims = [0, 2] in example above.
  other_dims = sorted(set(range(x_ndims)).difference(axis))
  # x_permed.shape = [a, c, b, d]
  perm = other_dims + list(axis) if right_end else list(axis) + other_dims
  x_permed = tf.transpose(a=x, perm=perm)

  if x.shape.is_fully_defined():
    x_shape = x.shape.as_list()
    # other_shape = [a, c], end_shape = [b * d]
    other_shape = [x_shape[i] for i in other_dims]
    end_shape = [np.prod([x_shape[i] for i in axis])]
    full_shape = (
        other_shape + end_shape if right_end else end_shape + other_shape)
  else:
    other_shape = tf.gather(tf.shape(input=x), other_dims)
    full_shape = tf.concat(
        [other_shape, [-1]] if right_end else [[-1], other_shape], axis=0)
  return tf.reshape(x_permed, shape=full_shape)

def _move_dims_to_flat_end ( x , axis , x_ndims , right_end = True ) : if not axis : return x other_dims = sorted ( set ( range ( x_ndims ) ) . difference ( axis ) ) perm = other_dims + list ( axis ) if right_end else list ( axis ) + other_dims x_permed = tf . transpose ( a = x , perm = perm ) if x . shape . is_fully_defined ( ) : x_shape = x . shape . as_list ( ) other_shape = [ x_shape [ i ] for i in other_dims ] end_shape = [ np . prod ( [ x_shape [ i ] for i in axis ] ) ] full_shape = ( other_shape + end_shape if right_end else end_shape + other_shape ) else : other_shape = tf . gather ( tf . shape ( input = x ) , other_dims ) full_shape = tf . concat ( [ other_shape , [ - 1 ] ] if right_end else [ [ - 1 ] , other_shape ] , axis = 0 ) return tf . reshape ( x_permed , shape = full_shape )

Move dims corresponding to axis in x to the end then flatten .


def _sort_tensor(tensor):
  """Use `top_k` to sort a `Tensor` along the last dimension."""
  sorted_, _ = tf.nn.top_k(tensor, k=tf.shape(input=tensor)[-1])
  sorted_.set_shape(tensor.shape)
  return sorted_

def _sort_tensor ( tensor ) : sorted_ , _ = tf . nn . top_k ( tensor , k = tf . shape ( input = tensor ) [ - 1 ] ) sorted_ . set_shape ( tensor . shape ) return sorted_

Use top_k to sort a Tensor along the last dimension .


def make_component_state_space_models(self,
                                        num_timesteps,
                                        param_vals,
                                        initial_step=0):
    """Build an ordered list of Distribution instances for component models.

    Args:
      num_timesteps: Python `int` number of timesteps to model.
      param_vals: a list of `Tensor` parameter values in order corresponding to
        `self.parameters`, or a dict mapping from parameter names to values.
      initial_step: optional `int` specifying the initial timestep to model.
        This is relevant when the model contains time-varying components,
        e.g., holidays or seasonality.

    Returns:
      component_ssms: a Python list of `LinearGaussianStateSpaceModel`
        Distribution objects, in order corresponding to `self.components`.
    """

    with tf.compat.v1.name_scope('make_component_state_space_models'):

      # List the model parameters in canonical order
      param_map = self._canonicalize_param_vals_as_map(param_vals)
      param_vals_list = [param_map[p.name] for p in self.parameters]

      # Build SSMs for each component model. We process the components in
      # canonical order, extracting the parameters for each component from the
      # (ordered) list of parameters.
      remaining_param_vals = param_vals_list[1:]
      component_ssms = []
      for component in self.components:
        num_parameters = len(component.parameters)
        component_param_vals = remaining_param_vals[:num_parameters]
        remaining_param_vals = remaining_param_vals[num_parameters:]

        component_ssms.append(
            component.make_state_space_model(
                num_timesteps,
                param_vals=component_param_vals,
                initial_step=initial_step))

    return component_ssms

def make_component_state_space_models ( self , num_timesteps , param_vals , initial_step = 0 ) : with tf . compat . v1 . name_scope ( 'make_component_state_space_models' ) : param_map = self . _canonicalize_param_vals_as_map ( param_vals ) param_vals_list = [ param_map [ p . name ] for p in self . parameters ] remaining_param_vals = param_vals_list [ 1 : ] component_ssms = [ ] for component in self . components : num_parameters = len ( component . parameters ) component_param_vals = remaining_param_vals [ : num_parameters ] remaining_param_vals = remaining_param_vals [ num_parameters : ] component_ssms . append ( component . make_state_space_model ( num_timesteps , param_vals = component_param_vals , initial_step = initial_step ) ) return component_ssms

Build an ordered list of Distribution instances for component models .


def amari_alpha(logu, alpha=1., self_normalized=False, name=None):
  """The Amari-alpha Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the Amari-alpha Csiszar-function is:

  ```none
  f(u) = { -log(u) + (u - 1),     alpha = 0
         { u log(u) - (u - 1),    alpha = 1
         { [(u**alpha - 1) - alpha (u - 1)] / (alpha (alpha - 1)),    otherwise
  ```

  When `self_normalized = False` the `(u - 1)` terms are omitted.

  Warning: when `alpha != 0` and/or `self_normalized = True` this function makes
  non-log-space calculations and may therefore be numerically unstable for
  `|logu| >> 0`.

  For more information, see:
    A. Cichocki and S. Amari. "Families of Alpha-Beta-and GammaDivergences:
    Flexible and Robust Measures of Similarities." Entropy, vol. 12, no. 6, pp.
    1532-1568, 2010.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    alpha: `float`-like Python scalar. (See Mathematical Details for meaning.)
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    amari_alpha_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.

  Raises:
    TypeError: if `alpha` is `None` or a `Tensor`.
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """
  with tf.compat.v1.name_scope(name, "amari_alpha", [logu]):
    if alpha is None or tf.is_tensor(alpha):
      raise TypeError("`alpha` cannot be `None` or `Tensor` type.")
    if (self_normalized is None or tf.is_tensor(self_normalized)):
      raise TypeError("`self_normalized` cannot be `None` or `Tensor` type.")

    logu = tf.convert_to_tensor(value=logu, name="logu")

    if alpha == 0.:
      f = -logu
    elif alpha == 1.:
      f = tf.exp(logu) * logu
    else:
      f = tf.math.expm1(alpha * logu) / (alpha * (alpha - 1.))

    if not self_normalized:
      return f

    if alpha == 0.:
      return f + tf.math.expm1(logu)
    elif alpha == 1.:
      return f - tf.math.expm1(logu)
    else:
      return f - tf.math.expm1(logu) / (alpha - 1.)

The Amari - alpha Csiszar - function in log - space .


def kl_reverse(logu, self_normalized=False, name=None):
  """The reverse Kullback-Leibler Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the KL-reverse Csiszar-function is:

  ```none
  f(u) = -log(u) + (u - 1)
  ```

  When `self_normalized = False` the `(u - 1)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[q, p]
  ```

  The KL is "reverse" because in maximum likelihood we think of minimizing `q`
  as in `KL[p, q]`.

  Warning: when self_normalized = True` this function makes non-log-space
  calculations and may therefore be numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    kl_reverse_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.

  Raises:
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """

  with tf.compat.v1.name_scope(name, "kl_reverse", [logu]):
    return amari_alpha(logu, alpha=0., self_normalized=self_normalized)

def kl_reverse ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "kl_reverse" , [ logu ] ) : return amari_alpha ( logu , alpha = 0. , self_normalized = self_normalized )

The reverse Kullback - Leibler Csiszar - function in log - space .


def jensen_shannon(logu, self_normalized=False, name=None):
  """The Jensen-Shannon Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the Jensen-Shannon Csiszar-function is:

  ```none
  f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2)
  ```

  When `self_normalized = False` the `(u + 1) log(2)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[p, m] + KL[q, m]
  m(x) = 0.5 p(x) + 0.5 q(x)
  ```

  In a sense, this divergence is the "reverse" of the Arithmetic-Geometric
  f-Divergence.

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  For more information, see:
    Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans.
    Inf. Th., 37, 145-151, 1991.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    jensen_shannon_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "jensen_shannon", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    npdt = logu.dtype.as_numpy_dtype
    y = tf.nn.softplus(logu)
    if self_normalized:
      y -= np.log(2).astype(npdt)
    return tf.exp(logu) * logu - (1. + tf.exp(logu)) * y

def jensen_shannon ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "jensen_shannon" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) npdt = logu . dtype . as_numpy_dtype y = tf . nn . softplus ( logu ) if self_normalized : y -= np . log ( 2 ) . astype ( npdt ) return tf . exp ( logu ) * logu - ( 1. + tf . exp ( logu ) ) * y

The Jensen - Shannon Csiszar - function in log - space .


def pearson(logu, name=None):
  """The Pearson Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Pearson Csiszar-function is:

  ```none
  f(u) = (u - 1)**2
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    pearson_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "pearson", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return tf.square(tf.math.expm1(logu))

def pearson ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "pearson" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return tf . square ( tf . math . expm1 ( logu ) )

The Pearson Csiszar - function in log - space .


def squared_hellinger(logu, name=None):
  """The Squared-Hellinger Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Squared-Hellinger Csiszar-function is:

  ```none
  f(u) = (sqrt(u) - 1)**2
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    squared_hellinger_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "squared_hellinger", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return pearson(0.5 * logu)

def squared_hellinger ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "squared_hellinger" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return pearson ( 0.5 * logu )

The Squared - Hellinger Csiszar - function in log - space .


def triangular(logu, name=None):
  """The Triangular Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Triangular Csiszar-function is:

  ```none
  f(u) = (u - 1)**2 / (1 + u)
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    triangular_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "triangular", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return pearson(logu) / (1. + tf.exp(logu))

def triangular ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "triangular" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return pearson ( logu ) / ( 1. + tf . exp ( logu ) )

The Triangular Csiszar - function in log - space .


def t_power(logu, t, self_normalized=False, name=None):
  """The T-Power Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True` the T-Power Csiszar-function is:

  ```none
  f(u) = s [ u**t - 1 - t(u - 1) ]
  s = { -1   0 < t < 1
      { +1   otherwise
  ```

  When `self_normalized = False` the `- t(u - 1)` term is omitted.

  This is similar to the `amari_alpha` Csiszar-function, with the associated
  divergence being the same up to factors depending only on `t`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    t:  `Tensor` of same `dtype` as `logu` and broadcastable shape.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    t_power_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """
  with tf.compat.v1.name_scope(name, "t_power", [logu, t]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    t = tf.convert_to_tensor(value=t, dtype=logu.dtype.base_dtype, name="t")
    fu = tf.math.expm1(t * logu)
    if self_normalized:
      fu -= t * tf.math.expm1(logu)
    fu *= tf.where(tf.logical_and(0. < t, t < 1.),
                   -tf.ones_like(t),
                   tf.ones_like(t))
    return fu

def t_power ( logu , t , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "t_power" , [ logu , t ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) t = tf . convert_to_tensor ( value = t , dtype = logu . dtype . base_dtype , name = "t" ) fu = tf . math . expm1 ( t * logu ) if self_normalized : fu -= t * tf . math . expm1 ( logu ) fu *= tf . where ( tf . logical_and ( 0. < t , t < 1. ) , - tf . ones_like ( t ) , tf . ones_like ( t ) ) return fu

The T - Power Csiszar - function in log - space .


def log1p_abs(logu, name=None):
  """The log1p-abs Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Log1p-Abs Csiszar-function is:

  ```none
  f(u) = u**(sign(u-1)) - 1
  ```

  This function is so-named because it was invented from the following recipe.
  Choose a convex function g such that g(0)=0 and solve for f:

  ```none
  log(1 + f(u)) = g(log(u)).
    <=>
  f(u) = exp(g(log(u))) - 1
  ```

  That is, the graph is identically `g` when y-axis is `log1p`-domain and x-axis
  is `log`-domain.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log1p_abs_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "log1p_abs", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return tf.math.expm1(tf.abs(logu))

def log1p_abs ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "log1p_abs" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return tf . math . expm1 ( tf . abs ( logu ) )

The log1p - abs Csiszar - function in log - space .


def jeffreys(logu, name=None):
  """The Jeffreys Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Jeffreys Csiszar-function is:

  ```none
  f(u) = 0.5 ( u log(u) - log(u) )
       = 0.5 kl_forward + 0.5 kl_reverse
       = symmetrized_csiszar_function(kl_reverse)
       = symmetrized_csiszar_function(kl_forward)
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    jeffreys_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "jeffreys", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return 0.5 * tf.math.expm1(logu) * logu

def jeffreys ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "jeffreys" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return 0.5 * tf . math . expm1 ( logu ) * logu

The Jeffreys Csiszar - function in log - space .


def modified_gan(logu, self_normalized=False, name=None):
  """The Modified-GAN Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True` the modified-GAN (Generative/Adversarial
  Network) Csiszar-function is:

  ```none
  f(u) = log(1 + u) - log(u) + 0.5 (u - 1)
  ```

  When `self_normalized = False` the `0.5 (u - 1)` is omitted.

  The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with
  `self_normalized = False`).

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "chi_square", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    y = tf.nn.softplus(logu) - logu
    if self_normalized:
      y += 0.5 * tf.math.expm1(logu)
    return y

def modified_gan ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "chi_square" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) y = tf . nn . softplus ( logu ) - logu if self_normalized : y += 0.5 * tf . math . expm1 ( logu ) return y

The Modified - GAN Csiszar - function in log - space .


def dual_csiszar_function(logu, csiszar_function, name=None):
  """Calculates the dual Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Csiszar-dual is defined as:

  ```none
  f^*(u) = u f(1 / u)
  ```

  where `f` is some other Csiszar-function.

  For example, the dual of `kl_reverse` is `kl_forward`, i.e.,

  ```none
  f(u) = -log(u)
  f^*(u) = u f(1 / u) = -u log(1 / u) = u log(u)
  ```

  The dual of the dual is the original function:

  ```none
  f^**(u) = {u f(1/u)}^*(u) = u (1/u) f(1/(1/u)) = f(u)
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    csiszar_function: Python `callable` representing a Csiszar-function over
      log-domain.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    dual_f_of_u: `float`-like `Tensor` of the result of calculating the dual of
      `f` at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "dual_csiszar_function", [logu]):
    return tf.exp(logu) * csiszar_function(-logu)

def dual_csiszar_function ( logu , csiszar_function , name = None ) : with tf . compat . v1 . name_scope ( name , "dual_csiszar_function" , [ logu ] ) : return tf . exp ( logu ) * csiszar_function ( - logu )

Calculates the dual Csiszar - function in log - space .


def symmetrized_csiszar_function(logu, csiszar_function, name=None):
  """Symmetrizes a Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The symmetrized Csiszar-function is defined as:

  ```none
  f_g(u) = 0.5 g(u) + 0.5 u g (1 / u)
  ```

  where `g` is some other Csiszar-function.

  We say the function is "symmetrized" because:

  ```none
  D_{f_g}[p, q] = D_{f_g}[q, p]
  ```

  for all `p << >> q` (i.e., `support(p) = support(q)`).

  There exists alternatives for symmetrizing a Csiszar-function. For example,

  ```none
  f_g(u) = max(f(u), f^*(u)),
  ```

  where `f^*` is the dual Csiszar-function, also implies a symmetric
  f-Divergence.

  Example:

  When either of the following functions are symmetrized, we obtain the
  Jensen-Shannon Csiszar-function, i.e.,

  ```none
  g(u) = -log(u) - (1 + u) log((1 + u) / 2) + u - 1
  h(u) = log(4) + 2 u log(u / (1 + u))
  ```

  implies,

  ```none
  f_g(u) = f_h(u) = u log(u) - (1 + u) log((1 + u) / 2)
         = jensen_shannon(log(u)).
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    csiszar_function: Python `callable` representing a Csiszar-function over
      log-domain.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    symmetrized_g_of_u: `float`-like `Tensor` of the result of applying the
      symmetrization of `g` evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "symmetrized_csiszar_function", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return 0.5 * (csiszar_function(logu)
                  + dual_csiszar_function(logu, csiszar_function))

def symmetrized_csiszar_function ( logu , csiszar_function , name = None ) : with tf . compat . v1 . name_scope ( name , "symmetrized_csiszar_function" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return 0.5 * ( csiszar_function ( logu ) + dual_csiszar_function ( logu , csiszar_function ) )

Symmetrizes a Csiszar - function in log - space .


def monte_carlo_csiszar_f_divergence(
    f,
    p_log_prob,
    q,
    num_draws,
    use_reparametrization=None,
    seed=None,
    name=None):
  """Monte-Carlo approximation of the Csiszar f-Divergence.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Csiszar f-Divergence for Csiszar-function f is given by:

  ```none
  D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ]
                  ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ),
                             where x_j ~iid q(X)
  ```

  Tricks: Reparameterization and Score-Gradient

  When q is "reparameterized", i.e., a diffeomorphic transformation of a
  parameterless distribution (e.g.,
  `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
  expectation, i.e.,
  `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}`
  and `s_i = f(x_i), x_i ~iid q(X)`.

  However, if q is not reparameterized, TensorFlow's gradient will be incorrect
  since the chain-rule stops at samples of unreparameterized distributions. In
  this circumstance using the Score-Gradient trick results in an unbiased
  gradient, i.e.,

  ```none
  grad[ E_q[f(X)] ]
  = grad[ int dx q(x) f(x) ]
  = int dx grad[ q(x) f(x) ]
  = int dx [ q'(x) f(x) + q(x) f'(x) ]
  = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ]
  = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ]
  = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ]
  ```

  Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is
  usually preferable to set `use_reparametrization = True`.

  Example Application:

  The Csiszar f-Divergence is a useful framework for variational inference.
  I.e., observe that,

  ```none
  f(p(x)) =  f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] )
          <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ]
          := D_f[p(x, Z), q(Z | x)]
  ```

  The inequality follows from the fact that the "perspective" of `f`, i.e.,
  `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and
  `t` is a real. Since the above framework includes the popular Evidence Lower
  BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework
  "Evidence Divergence Bound Optimization" (EDBO).

  Args:
    f: Python `callable` representing a Csiszar-function in log-space, i.e.,
      takes `p_log_prob(q_samples) - q.log_prob(q_samples)`.
    p_log_prob: Python `callable` taking (a batch of) samples from `q` and
      returning the natural-log of the probability under distribution `p`.
      (In variational inference `p` is the joint distribution.)
    q: `tf.Distribution`-like instance; must implement:
      `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`.
      (In variational inference `q` is the approximate posterior distribution.)
    num_draws: Integer scalar number of draws used to approximate the
      f-Divergence expectation.
    use_reparametrization: Python `bool`. When `None` (the default),
      automatically set to:
      `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`.
      When `True` uses the standard Monte-Carlo average. When `False` uses the
      score-gradient trick. (See above for details.)  When `False`, consider
      using `csiszar_vimco`.
    seed: Python `int` seed for `q.sample`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo
      approximation of the Csiszar f-Divergence.

  Raises:
    ValueError: if `q` is not a reparameterized distribution and
      `use_reparametrization = True`. A distribution `q` is said to be
      "reparameterized" when its samples are generated by transforming the
      samples of another distribution which does not depend on the
      parameterization of `q`. This property ensures the gradient (with respect
      to parameters) is valid.
    TypeError: if `p_log_prob` is not a Python `callable`.
  """
  reparameterization_types = tf.nest.flatten(q.reparameterization_type)
  with tf.compat.v1.name_scope(name, "monte_carlo_csiszar_f_divergence",
                               [num_draws]):
    if use_reparametrization is None:
      use_reparametrization = all(
          reparameterization_type == tfd.FULLY_REPARAMETERIZED
          for reparameterization_type in reparameterization_types)
    elif (use_reparametrization and
          any(reparameterization_type != tfd.FULLY_REPARAMETERIZED
              for reparameterization_type in reparameterization_types)):
      # TODO(jvdillon): Consider only raising an exception if the gradient is
      # requested.
      raise ValueError(
          "Distribution `q` must be reparameterized, i.e., a diffeomorphic "
          "transformation of a parameterless distribution. (Otherwise this "
          "function has a biased gradient.)")
    if not callable(p_log_prob):
      raise TypeError("`p_log_prob` must be a Python `callable` function.")
    return monte_carlo.expectation(
        f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)),
        samples=q.sample(num_draws, seed=seed),
        log_prob=q.log_prob,  # Only used if use_reparametrization=False.
        use_reparametrization=use_reparametrization)

Monte - Carlo approximation of the Csiszar f - Divergence .