class_docstr: Posterior predictive energy distance [1,2] with optional Bayesian
  regularization by the prior.
  
  Let `p(x,z)=p(z) p(x|z)` be the model, `q(z|x)` be the guide.  Then given
  data `x` and drawing an iid pair of samples :math:`(Z,X)` and
  :math:`(Z',X')` (where `Z` is latent and `X` is the posterior predictive),
  
  .. math ::
  
      & Z \sim q(z|x); \quad X \sim p(x|Z) \\
      & Z' \sim q(z|x); \quad X' \sim p(x|Z') \\
      & loss = \mathbb E_X \|X-x\|^\beta
             - \frac 1 2 \mathbb E_{X,X'}\|X-X'\|^\beta
             - \lambda \mathbb E_Z \log p(Z)
  
  This is a likelihood-free inference algorithm, and can be used for
  likelihoods without tractable density functions. The :math:`\beta` energy
  distance is a robust loss functions, and is well defined for any
  distribution with finite fractional moment :math:`\mathbb E[\|X\|^\beta]`.
  
  This requires static model structure, a fully reparametrized guide, and
  reparametrized likelihood distributions in the model. Model latent
  distributions may be non-reparametrized.
  
  **References**
  
  [1] Gabor J. Szekely, Maria L. Rizzo (2003)
      Energy Statistics: A Class of Statistics Based on Distances.
  [2] Tilmann Gneiting, Adrian E. Raftery (2007)
      Strictly Proper Scoring Rules, Prediction, and Estimation.
      https://www.stat.washington.edu/raftery/Research/PDF/Gneiting2007jasa.pdf
  
  :param float beta: Exponent :math:`\beta` from [1,2]. The loss function is
      strictly proper for distributions with finite :math:`beta`-absolute moment
      :math:`E[\|X\|^\beta]`. Thus for heavy tailed distributions ``beta`` should
      be small, e.g. for ``Cauchy`` distributions, :math:`\beta<1` is strictly
      proper. Defaults to 1. Must be in the open interval (0,2).
  :param float prior_scale: Nonnegative scale for prior regularization.
      Model parameters are trained only if this is positive.
      If zero (default), then model log densities will not be computed
      (guide log densities are never computed).
  :param int num_particles: The number of particles/samples used to form the
      gradient estimators. Must be at least 2.
  :param int max_plate_nesting: Optional bound on max number of nested
      :func:`pyro.plate` contexts. If omitted, this will guess a valid value
      by running the (model,guide) pair once.
class_docstr_tok: ["Posterior","predictive","energy","distance","[","1,2","]","with","optional","Bayesian","regularization","by","the","prior",".","Let","`","p","(","x",",","z",")","=p","(","z",")","p","(","x|z",")","`","be","the","model",",","`","q","(","z|x",")","`","be","the","guide",".","Then","given","data","`","x","`","and","drawing","an","iid","pair","of","samples",":","math",":","`","(","Z",",","X",")","`","and",":","math",":","`","(","Z","'",",","X","'",")","`","(","where","`","Z","`","is","latent","and","`","X","`","is","the","posterior","predictive",")",",","..","math",":",":","&","Z","\\sim","q","(","z|x",")",";","\\quad","X","\\sim","p","(","x|Z",")","\\\\","&","Z","'","\\sim","q","(","z|x",")",";","\\quad","X","'","\\sim","p","(","x|Z","'",")","\\\\","&","loss","=","\\mathbb","E_X","\\|X-x\\|^\\beta","-","\\frac","1","2","\\mathbb","E_","{","X",",","X","'","}","\\|X-X'\\|^\\beta","-","\\lambda","\\mathbb","E_Z","\\log","p","(","Z",")","This","is","a","likelihood-free","inference","algorithm",",","and","can","be","used","for","likelihoods","without","tractable","density","functions",".","The",":","math",":","`","\\beta","`","energy","distance","is","a","robust","loss","functions",",","and","is","well","defined","for","any","distribution","with","finite","fractional","moment",":","math",":","`","\\mathbb","E","[","\\|X\\|^\\beta","]","`",".","This","requires","static","model","structure",",","a","fully","reparametrized","guide",",","and","reparametrized","likelihood","distributions","in","the","model",".","Model","latent","distributions","may","be","non-reparametrized",".","*","*","References","*","*","[","1","]","Gabor","J.","Szekely",",","Maria","L.","Rizzo","(","2003",")","Energy","Statistics",":","A","Class","of","Statistics","Based","on","Distances",".","[","2","]","Tilmann","Gneiting",",","Adrian","E.","Raftery","(","2007",")","Strictly","Proper","Scoring","Rules",",","Prediction",",","and","Estimation",".","https",":","\/\/www.stat.washington.edu\/raftery\/Research\/PDF\/Gneiting2007jasa.pdf",":","param","float","beta",":","Exponent",":","math",":","`","\\beta","`","from","[","1,2","]",".","The","loss","function","is","strictly","proper","for","distributions","with","finite",":","math",":","`","beta","`","-absolute","moment",":","math",":","`","E","[","\\|X\\|^\\beta","]","`",".","Thus","for","heavy","tailed","distributions","``","beta","``","should","be","small",",","e.g",".","for","``","Cauchy","``","distributions",",",":","math",":","`","\\beta","<","1","`","is","strictly","proper",".","Defaults","to","1",".","Must","be","in","the","open","interval","(","0,2",")",".",":","param","float","prior_scale",":","Nonnegative","scale","for","prior","regularization",".","Model","parameters","are","trained","only","if","this","is","positive",".","If","zero","(","default",")",",","then","model","log","densities","will","not","be","computed","(","guide","log","densities","are","never","computed",")",".",":","param","int","num_particles",":","The","number","of","particles\/samples","used","to","form","the","gradient","estimators",".","Must","be","at","least","2.",":","param","int","max_plate_nesting",":","Optional","bound","on","max","number","of","nested",":","func",":","`","pyro.plate","`","contexts",".","If","omitted",",","this","will","guess","a","valid","value","by","running","the","(","model",",","guide",")","pair","once","."]
code: def _get_traces(self, model, guide, args, kwargs):
      if self.max_plate_nesting == float("inf"):
          with validation_enabled(
              False
          ):  # Avoid calling .log_prob() when undefined.
              # TODO factor this out as a stand-alone helper.
              ELBO._guess_max_plate_nesting(
                  self, model, guide, args, kwargs
              )
      vectorize = pyro.plate(
          "num_particles_vectorized",
          self.num_particles,
          dim=-self.max_plate_nesting,
      )
  
      # Trace the guide as in ELBO.
      with poutine.trace() as tr, vectorize:
          guide(*args, **kwargs)
      guide_trace = tr.trace
  
      # Trace the model, drawing posterior predictive samples.
      with poutine.trace() as tr, poutine.uncondition():
          with poutine.replay(trace=guide_trace), vectorize:
              model(*args, **kwargs)
      model_trace = tr.trace
      for site in model_trace.nodes.values():
          if site["type"] == "sample" and site["infer"].get(
              "was_observed", False
          ):
              site["is_observed"] = True
      if is_validation_enabled():
          check_model_guide_match(
              model_trace, guide_trace, self.max_plate_nesting
          )
  
      guide_trace = prune_subsample_sites(guide_trace)
      model_trace = prune_subsample_sites(model_trace)
      if is_validation_enabled():
          for site in guide_trace.nodes.values():
              if site["type"] == "sample":
                  warn_if_nan(site["value"], site["name"])
                  if not getattr(site["fn"], "has_rsample", False):
                      raise ValueError(
                          "EnergyDistance requires fully reparametrized guides"
                      )
          for trace in model_trace.nodes.values():
              if site["type"] == "sample":
                  if site["is_observed"]:
                      warn_if_nan(site["value"], site["name"])
                      if not getattr(site["fn"], "has_rsample", False):
                          raise ValueError(
                              "EnergyDistance requires reparametrized likelihoods"
                          )
  
      if self.prior_scale > 0:
          model_trace.compute_log_prob(
              site_filter=lambda name, site: not site["is_observed"]
          )
          if is_validation_enabled():
              for site in model_trace.nodes.values():
                  if site["type"] == "sample":
                      if not site["is_observed"]:
                          check_site_shape(site, self.max_plate_nesting)
  
      return guide_trace, model_trace
  
file_path: pyro/pyro/infer/energy_distance.py
imports: import operator
  from collections import OrderedDict
  from functools import reduce
  import torch
  import pyro
  import pyro.poutine
  from pyro.distributions.util import scale_and_mask
  from pyro.infer.elbo import ELBO
  from pyro.infer.util import is_validation_enabled, validation_enabled
  from pyro.poutine.util import prune_subsample_sites
  from pyro.util import check_model_guide_match, check_site_shape, warn_if_nan
label_desc: Use image node_id 3 for calling the EnergyDistance obj's underlying member method code with example usage: obj._get_traces(model, guide, args, kwargs) and returns: guide_trace,  model_trace