class_docstr: Models the distribution of finitely many importance-reweighted samples.
  
  This wrapper adapts a proposal distribution towards a target density using
  [importance sampling](https://en.wikipedia.org/wiki/Importance_sampling).
  Given a proposal `q`, a target density `p` (which may be unnormalized), and
  an integer `importance_sample_size`, it models the result of the following
  sampling process:
  
  1. Draw `importance_sample_size` samples `z[k] ~ q` from the proposal.
  2. Compute an importance weight `w[k] = p(z[k]) / q(z[k])` for each sample.
  3. Return a sample `z[k*]` selected with probability proportional to
     the importance weights, i.e., with `k* ~ Categorical(probs=w/sum(w))`.
  
  In the limit where `importance_sample_size -> inf`, the result `z[k*]` of this
  procedure would be distributed according to the target density `p`. On the
  other hand, if `importance_sample_size == 1`, then the reweighting has no
  effect and the result `z[k*]` is simply a sample from `q`. Finite values
  of `importance_sample_size` describe distributions that are intermediate
  between `p` and `q`.
  
  This distribution may also be understood as an explicit representation of the
  surrogate posterior that is implicitly assumed by importance-weighted
  variational objectives. [1, 2]
  
  #### Examples
  
  This distribution can be used directly for posterior inference via importance
  sampling:
  
  ```python
  tfd = tfp.distributions
  tfed = tfp.experimental.distributions
  
  def target_log_prob_fn(x):
    prior = tfd.Normal(loc=0., scale=1.).log_prob(x)
    # Multimodal likelihood.
    likelihood = tf.reduce_logsumexp(
      tfd.Normal(loc=x, scale=0.1).log_prob([-1., 1.]))
    return prior + likelihood
  
  # Use importance sampling to infer an approximate posterior.
  approximate_posterior = tfed.ImportanceResample(
    proposal_distribution=tfd.Normal(loc=0., scale=2.),
    target_log_prob_fn=target_log_prob_fn,
    importance_sample_size=100)
  ```
  
  We can estimate posterior expectations directly using an importance-weighted
  sum of proposal samples:
  
  ```python
  # Directly compute expectations under the posterior via importance weights.
  posterior_mean = approximate_posterior.self_normalized_expectation(
    lambda x: x, importance_sample_size=1000)
  posterior_variance = approximate_posterior.self_normalized_expectation(
    lambda x: (x - posterior_mean)**2, importance_sample_size=1000)
  ```
  
  Alternately, the same expectations can be estimated from explicit (unweighted)
  samples. Note that sampling may be expensive because it performs resampling
  internally. For example, to produce `sample_size` samples requires first
  proposing values of shape `[sample_size, importance_sample_size]`
  (`[1000, 100]` in the code below) and then resampling down to `[sample_size]`,
  throwing most of the proposals away. For this reason you should prefer calling
  `self_normalized_expectation` over naive sampling to compute expectations.
  
  ```python
  posterior_samples = approximate_posterior.sample(1000)
  posterior_mean_inefficient = tf.reduce_mean(posterior_samples)
  posterior_variance_inefficient = tf.math.reduce_variance(posterior_samples)
  
  # Calling `self_normalized_expectation` allows for a much lower `sample_size`
  # because it uses the full set of `importance_sample_size` proposal samples to
  # approximate the expectation at each of the `sample_size` Monte Carlo
  # evaluations. This is formalized in Eq. 9 of [3].
  posterior_mean_efficient = approximate_posterior.self_normalized_expectation(
    lambda x: x, sample_size=10)
  posterior_variance_efficient = (
    approximate_posterior.self_normalized_expectation(
      lambda x: (x - posterior_mean_efficient)**2, sample_size=10))
  ```
  
  The posterior (log-)density cannot be computed directly, but may be
  stochastically approximated. The `prob` and `log_prob` methods accept
  arguments `seed` and `sample_size` to control the variance of the
  approximation.
  
  ```python
  # Plot the posterior density.
  from matplotlib import pylab as plt
  xs = tf.linspace(-3., 3., 101)
  probs = approximate_posterior.prob(xs, sample_size=10, seed=(42, 42))
  plt.plot(xs, probs)
  ```
  
  #### Connections to importance-weighted variational inference
  
  Optimizing an importance-weighted variational bound provides a natural
  approach to choose a proposal distribution for importance sampling.
  Importance-weighted bounds are available directly in TFP via the
  `importance_sample_size` argument to `tfp.vi.monte_carlo_variational_loss`
  and `tfp.vi.fit_surrogate_posterior`. For example, we might improve on the
  example above by replacing the fixed proposal distribution with a learned
  proposal:
  
  ```python
  proposal_distribution = tfp.experimental.util.make_trainable(tfd.Normal)
  importance_sample_size = 100
  importance_weighted_losses = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn,
    surrogate_posterior=proposal_distribution,
    optimizer=tf_keras.optimizers.Adam(0.1),
    num_steps=200,
    importance_sample_size=importance_sample_size)
  approximate_posterior = tfed.ImportanceResample(
    proposal_distribution=proposal_distribution,
    target_log_prob_fn=target_log_prob_fn,
    importance_sample_size=importance_sample_size)
  ```
  
  Note that although the importance-resampled `approximate_posterior` serves
  ultimately as the surrogate posterior, only the bare proposal distribution
  is passed as the `surrogate_posterior` argument to `fit_surrogate_posterior`.
  This is because the `importance_sample_size` argument tells
  `fit_surrogate_posterior` to compute an importance-weighted bound directly
  from the proposal distribution. Mathematically, it would be equivalent to omit
  the `importance_sample_size` argument and instead pass an `ImportanceResample`
  distribution as the surrogate posterior:
  
  ```python
  equivalent_but_less_efficient_losses = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn,
    surrogate_posterior=tfed.ImportanceResample(
      proposal_distribution=proposal_distribution,
      target_log_prob_fn=target_log_prob_fn,
      importance_sample_size=importance_sample_size),
    optimizer=tf_keras.optimizers.Adam(0.1),
    num_steps=200)
  ```
  
  but this approach is not recommended, because it performs redundant
  evaluations of the `target_log_prob_fn` compared to the direct bound shown
  above.
  
  #### References
  
  [1] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance Weighted
      Autoencoders. In _International Conference on Learning
      Representations_, 2016. https://arxiv.org/abs/1509.00519
  [2] Chris Cremer, Quaid Morris, and David Duvenaud. Reinterpreting
      Importance-Weighted Autoencoders. In _International Conference on Learning
      Representations_, Workshop track, 2017. https://arxiv.org/abs/1704.02916
  [3] Justin Domke, Daniel Sheldon. Importance Weighting and Variational
      Inference. In _Neural Information Processing Systems (NIPS)_, 2018.
      https://arxiv.org/abs/1808.09034
class_docstr_tok: ["Models","the","distribution","of","finitely","many","importance-reweighted","samples",".","This","wrapper","adapts","a","proposal","distribution","towards","a","target","density","using","[","importance","sampling","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Importance_sampling",")",".","Given","a","proposal","`","q","`",",","a","target","density","`","p","`","(","which","may","be","unnormalized",")",",","and","an","integer","`","importance_sample_size","`",",","it","models","the","result","of","the","following","sampling","process",":","1",".","Draw","`","importance_sample_size","`","samples","`","z","[","k","]","~","q","`","from","the","proposal",".","2",".","Compute","an","importance","weight","`","w","[","k","]","=","p","(","z","[","k","]",")","\/","q","(","z","[","k","]",")","`","for","each","sample",".","3",".","Return","a","sample","`","z","[","k","*","]","`","selected","with","probability","proportional","to","the","importance","weights",",","i.e.",",","with","`","k","*","~","Categorical","(","probs=w\/sum","(","w",")",")","`",".","In","the","limit","where","`","importance_sample_size","-",">","inf","`",",","the","result","`","z","[","k","*","]","`","of","this","procedure","would","be","distributed","according","to","the","target","density","`","p","`",".","On","the","other","hand",",","if","`","importance_sample_size","==","1","`",",","then","the","reweighting","has","no","effect","and","the","result","`","z","[","k","*","]","`","is","simply","a","sample","from","`","q","`",".","Finite","values","of","`","importance_sample_size","`","describe","distributions","that","are","intermediate","between","`","p","`","and","`","q","`",".","This","distribution","may","also","be","understood","as","an","explicit","representation","of","the","surrogate","posterior","that","is","implicitly","assumed","by","importance-weighted","variational","objectives",".","[","1",",","2","]","#","#","#","#","Examples","This","distribution","can","be","used","directly","for","posterior","inference","via","importance","sampling",":","``","`","python","tfd","=","tfp.distributions","tfed","=","tfp.experimental.distributions","def","target_log_prob_fn","(","x",")",":","prior","=","tfd.Normal","(","loc=0.",",","scale=1",".",")",".log_prob","(","x",")","#","Multimodal","likelihood",".","likelihood","=","tf.reduce_logsumexp","(","tfd.Normal","(","loc=x",",","scale=0.1",")",".log_prob","(","[","-1.",",","1",".","]",")",")","return","prior","+","likelihood","#","Use","importance","sampling","to","infer","an","approximate","posterior",".","approximate_posterior","=","tfed.ImportanceResample","(","proposal_distribution=tfd.Normal","(","loc=0.",",","scale=2",".",")",",","target_log_prob_fn=target_log_prob_fn",",","importance_sample_size=100",")","``","`","We","can","estimate","posterior","expectations","directly","using","an","importance-weighted","sum","of","proposal","samples",":","``","`","python","#","Directly","compute","expectations","under","the","posterior","via","importance","weights",".","posterior_mean","=","approximate_posterior.self_normalized_expectation","(","lambda","x",":","x",",","importance_sample_size=1000",")","posterior_variance","=","approximate_posterior.self_normalized_expectation","(","lambda","x",":","(","x","-","posterior_mean",")","*","*","2",",","importance_sample_size=1000",")","``","`","Alternately",",","the","same","expectations","can","be","estimated","from","explicit","(","unweighted",")","samples",".","Note","that","sampling","may","be","expensive","because","it","performs","resampling","internally",".","For","example",",","to","produce","`","sample_size","`","samples","requires","first","proposing","values","of","shape","`","[","sample_size",",","importance_sample_size","]","`","(","`","[","1000",",","100","]","`","in","the","code","below",")","and","then","resampling","down","to","`","[","sample_size","]","`",",","throwing","most","of","the","proposals","away",".","For","this","reason","you","should","prefer","calling","`","self_normalized_expectation","`","over","naive","sampling","to","compute","expectations",".","``","`","python","posterior_samples","=","approximate_posterior.sample","(","1000",")","posterior_mean_inefficient","=","tf.reduce_mean","(","posterior_samples",")","posterior_variance_inefficient","=","tf.math.reduce_variance","(","posterior_samples",")","#","Calling","`","self_normalized_expectation","`","allows","for","a","much","lower","`","sample_size","`","#","because","it","uses","the","full","set","of","`","importance_sample_size","`","proposal","samples","to","#","approximate","the","expectation","at","each","of","the","`","sample_size","`","Monte","Carlo","#","evaluations",".","This","is","formalized","in","Eq",".","9","of","[","3","]",".","posterior_mean_efficient","=","approximate_posterior.self_normalized_expectation","(","lambda","x",":","x",",","sample_size=10",")","posterior_variance_efficient","=","(","approximate_posterior.self_normalized_expectation","(","lambda","x",":","(","x","-","posterior_mean_efficient",")","*","*","2",",","sample_size=10",")",")","``","`","The","posterior","(","log-",")","density","can","not","be","computed","directly",",","but","may","be","stochastically","approximated",".","The","`","prob","`","and","`","log_prob","`","methods","accept","arguments","`","seed","`","and","`","sample_size","`","to","control","the","variance","of","the","approximation",".","``","`","python","#","Plot","the","posterior","density",".","from","matplotlib","import","pylab","as","plt","xs","=","tf.linspace","(","-3.",",","3.",",","101",")","probs","=","approximate_posterior.prob","(","xs",",","sample_size=10",",","seed=","(","42",",","42",")",")","plt.plot","(","xs",",","probs",")","``","`","#","#","#","#","Connections","to","importance-weighted","variational","inference","Optimizing","an","importance-weighted","variational","bound","provides","a","natural","approach","to","choose","a","proposal","distribution","for","importance","sampling",".","Importance-weighted","bounds","are","available","directly","in","TFP","via","the","`","importance_sample_size","`","argument","to","`","tfp.vi.monte_carlo_variational_loss","`","and","`","tfp.vi.fit_surrogate_posterior","`",".","For","example",",","we","might","improve","on","the","example","above","by","replacing","the","fixed","proposal","distribution","with","a","learned","proposal",":","``","`","python","proposal_distribution","=","tfp.experimental.util.make_trainable","(","tfd.Normal",")","importance_sample_size","=","100","importance_weighted_losses","=","tfp.vi.fit_surrogate_posterior","(","target_log_prob_fn",",","surrogate_posterior=proposal_distribution",",","optimizer=tf_keras.optimizers.Adam","(","0.1",")",",","num_steps=200",",","importance_sample_size=importance_sample_size",")","approximate_posterior","=","tfed.ImportanceResample","(","proposal_distribution=proposal_distribution",",","target_log_prob_fn=target_log_prob_fn",",","importance_sample_size=importance_sample_size",")","``","`","Note","that","although","the","importance-resampled","`","approximate_posterior","`","serves","ultimately","as","the","surrogate","posterior",",","only","the","bare","proposal","distribution","is","passed","as","the","`","surrogate_posterior","`","argument","to","`","fit_surrogate_posterior","`",".","This","is","because","the","`","importance_sample_size","`","argument","tells","`","fit_surrogate_posterior","`","to","compute","an","importance-weighted","bound","directly","from","the","proposal","distribution",".","Mathematically",",","it","would","be","equivalent","to","omit","the","`","importance_sample_size","`","argument","and","instead","pass","an","`","ImportanceResample","`","distribution","as","the","surrogate","posterior",":","``","`","python","equivalent_but_less_efficient_losses","=","tfp.vi.fit_surrogate_posterior","(","target_log_prob_fn",",","surrogate_posterior=tfed.ImportanceResample","(","proposal_distribution=proposal_distribution",",","target_log_prob_fn=target_log_prob_fn",",","importance_sample_size=importance_sample_size",")",",","optimizer=tf_keras.optimizers.Adam","(","0.1",")",",","num_steps=200",")","``","`","but","this","approach","is","not","recommended",",","because","it","performs","redundant","evaluations","of","the","`","target_log_prob_fn","`","compared","to","the","direct","bound","shown","above",".","#","#","#","#","References","[","1","]","Yuri","Burda",",","Roger","Grosse",",","and","Ruslan","Salakhutdinov",".","Importance","Weighted","Autoencoders",".","In","_International","Conference","on","Learning","Representations_",",","2016.","https",":","\/\/arxiv.org\/abs\/1509.00519","[","2","]","Chris","Cremer",",","Quaid","Morris",",","and","David","Duvenaud",".","Reinterpreting","Importance-Weighted","Autoencoders",".","In","_International","Conference","on","Learning","Representations_",",","Workshop","track",",","2017.","https",":","\/\/arxiv.org\/abs\/1704.02916","[","3","]","Justin","Domke",",","Daniel","Sheldon",".","Importance","Weighting","and","Variational","Inference",".","In","_Neural","Information","Processing","Systems","(","NIPS",")","_",",","2018.","https",":","\/\/arxiv.org\/abs\/1808.09034"]
code: def _propose_with_log_weights(self, sample_shape, seed=None):
      (
          particles,
          proposal_log_prob,
      ) = self.proposal_distribution.experimental_sample_and_log_prob(
          sample_shape, seed=seed
      )
      log_weights = (
          self._call_target_log_prob_fn(particles) - proposal_log_prob
      )
      # Ensure that the target fn didn't add batch dimensions.
      with tf.control_dependencies(
          self._check_weights_shape(
              log_weights, sample_shape=sample_shape
          )
      ):
          log_weights = tf.identity(log_weights)
  
      return particles, log_weights
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 15 for calling the ImportanceResample obj's underlying member method code with example usage: obj._propose_with_log_weights(sample_shape, seed) and returns: particles,  log_weights