class_docstr: Joint distribution parameterized by named distribution-making functions.
  
  This class provides automatic vectorization and alternative semantics for
  `tfd.JointDistributionNamed`, which in many cases allows for
  simplifications in the model specification.
  
  #### Automatic vectorization
  
  Auto-vectorized variants of JointDistribution allow the user to avoid
  explicitly annotating a model's vectorization semantics.
  When using manually-vectorized joint distributions, each operation in the
  model must account for the possibility of batch dimensions in Distributions
  and their samples. By contrast, auto-vectorized models need only describe
  a *single* sample from the joint distribution; any batch evaluation is
  automated using `tf.vectorized_map` as required. In many cases this
  allows for significant simplications. For example, the following
  manually-vectorized `tfd.JointDistributionNamed` model:
  
  ```python
  model = tfd.JointDistributionNamed({
    'x': tfd.Normal(0., tf.ones([3])),
    'y': tfd.Normal(0., 1.),
    'z': lambda x, y: tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)
  })
  ```
  
  can be written in auto-vectorized form as
  
  ```python
  model = tfd.JointDistributionNamedAutoBatched({
    'x': tfd.Normal(0., tf.ones([3])),
    'y': tfd.Normal(0., 1.),
    'z': lambda x, y: tfd.Normal(x[:2] + y, 1.)
  })
  ```
  
  in which we were able to avoid explicitly accounting for batch dimensions
  when indexing and slicing computed quantities in the third line.
  
  Note: auto-vectorization is still experimental and some TensorFlow ops may
  be unsupported. It can be disabled by setting `use_vectorized_map=False`.
  
  #### Alternative batch semantics
  
  This class also provides alternative semantics for specifying a batch of
  independent (non-identical) joint distributions.
  
  Instead of simply summing the `log_prob`s of component distributions
  (which may have different shapes), it first reduces the component `log_prob`s
  to ensure that `jd.log_prob(jd.sample())` always returns a scalar, unless
  `batch_ndims` is explicitly set to a nonzero value (in which case the result
  will have the corresponding tensor rank).
  
  The essential changes are:
  
  - An `event` of `JointDistributionNamedAutoBatched` is the dictionary of
    tensors produced by `.sample()`; thus, the `event_shape` is the
    dictionary containing the shapes of sampled tensors. These combine both
    the event and batch dimensions of the component distributions. By contrast,
    the event shape of a base `JointDistribution`s does not include batch
    dimensions of component distributions.
  - The `batch_shape` is a global property of the entire model, rather
    than a per-component property as in base `JointDistribution`s.
    The global batch shape must be a prefix of the batch shapes of
    each component; the length of this prefix is specified by an optional
    argument `batch_ndims`. If `batch_ndims` is not specified, the model has
    batch shape `[]`.
  
  #### Examples
  
  Consider the following generative model:
  
  ```
  e ~ Exponential(rate=[100,120])
  g ~ Gamma(concentration=e[0], rate=e[1])
  n ~ Normal(loc=0, scale=2.)
  m ~ Normal(loc=n, scale=g)
  for i = 1, ..., 12:
    x[i] ~ Bernoulli(logits=m)
  ```
  
  We can code this as:
  
  ```python
  tfd = tfp.distributions
  joint = tfd.JointDistributionNamedAutoBatched(dict(
      e=             tfd.Exponential(rate=[100, 120]),
      g=lambda    e: tfd.Gamma(concentration=e[0], rate=e[1]),
      n=             tfd.Normal(loc=0, scale=2.),
      m=lambda n, g: tfd.Normal(loc=n, scale=g),
      x=lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12),
  ))
  ```
  
  Notice the 1:1 correspondence between "math" and "code". In a standard
  `JointDistributionNamed`, we would have wrapped the first variable as
  `e = tfd.Independent(tfd.Exponential(rate=[100, 120]),
   reinterpreted_batch_ndims=1)` to specify that `log_prob` of the `Exponential`
  should be a scalar, summing over both dimensions. We would also have had to
  extend indices as `tfd.Gamma(concentration=e[..., 0], rate=e[..., 1])` to
  account for possible batch dimensions. Both of these behaviors are implicit
  in `JointDistributionNamedAutoBatched`.
class_docstr_tok: ["Joint","distribution","parameterized","by","named","distribution-making","functions",".","This","class","provides","automatic","vectorization","and","alternative","semantics","for","`","tfd.JointDistributionNamed","`",",","which","in","many","cases","allows","for","simplifications","in","the","model","specification",".","#","#","#","#","Automatic","vectorization","Auto-vectorized","variants","of","JointDistribution","allow","the","user","to","avoid","explicitly","annotating","a","model","'s","vectorization","semantics",".","When","using","manually-vectorized","joint","distributions",",","each","operation","in","the","model","must","account","for","the","possibility","of","batch","dimensions","in","Distributions","and","their","samples",".","By","contrast",",","auto-vectorized","models","need","only","describe","a","*","single","*","sample","from","the","joint","distribution",";","any","batch","evaluation","is","automated","using","`","tf.vectorized_map","`","as","required",".","In","many","cases","this","allows","for","significant","simplications",".","For","example",",","the","following","manually-vectorized","`","tfd.JointDistributionNamed","`","model",":","``","`","python","model","=","tfd.JointDistributionNamed","(","{","'","x","'",":","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","'","y","'",":","tfd.Normal","(","0.",",","1",".",")",",","'","z","'",":","lambda","x",",","y",":","tfd.Normal","(","x","[","...",",",":2","]","+","y","[","...",",","tf.newaxis","]",",","1",".",")","}",")","``","`","can","be","written","in","auto-vectorized","form","as","``","`","python","model","=","tfd.JointDistributionNamedAutoBatched","(","{","'","x","'",":","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","'","y","'",":","tfd.Normal","(","0.",",","1",".",")",",","'","z","'",":","lambda","x",",","y",":","tfd.Normal","(","x","[",":2","]","+","y",",","1",".",")","}",")","``","`","in","which","we","were","able","to","avoid","explicitly","accounting","for","batch","dimensions","when","indexing","and","slicing","computed","quantities","in","the","third","line",".","Note",":","auto-vectorization","is","still","experimental","and","some","TensorFlow","ops","may","be","unsupported",".","It","can","be","disabled","by","setting","`","use_vectorized_map=False","`",".","#","#","#","#","Alternative","batch","semantics","This","class","also","provides","alternative","semantics","for","specifying","a","batch","of","independent","(","non-identical",")","joint","distributions",".","Instead","of","simply","summing","the","`","log_prob","`","s","of","component","distributions","(","which","may","have","different","shapes",")",",","it","first","reduces","the","component","`","log_prob","`","s","to","ensure","that","`","jd.log_prob","(","jd.sample","(",")",")","`","always","returns","a","scalar",",","unless","`","batch_ndims","`","is","explicitly","set","to","a","nonzero","value","(","in","which","case","the","result","will","have","the","corresponding","tensor","rank",")",".","The","essential","changes","are",":","-","An","`","event","`","of","`","JointDistributionNamedAutoBatched","`","is","the","dictionary","of","tensors","produced","by","`",".sample","(",")","`",";","thus",",","the","`","event_shape","`","is","the","dictionary","containing","the","shapes","of","sampled","tensors",".","These","combine","both","the","event","and","batch","dimensions","of","the","component","distributions",".","By","contrast",",","the","event","shape","of","a","base","`","JointDistribution","`","s","does","not","include","batch","dimensions","of","component","distributions",".","-","The","`","batch_shape","`","is","a","global","property","of","the","entire","model",",","rather","than","a","per-component","property","as","in","base","`","JointDistribution","`","s",".","The","global","batch","shape","must","be","a","prefix","of","the","batch","shapes","of","each","component",";","the","length","of","this","prefix","is","specified","by","an","optional","argument","`","batch_ndims","`",".","If","`","batch_ndims","`","is","not","specified",",","the","model","has","batch","shape","`","[","]","`",".","#","#","#","#","Examples","Consider","the","following","generative","model",":","``","`","e","~","Exponential","(","rate=","[","100,120","]",")","g","~","Gamma","(","concentration=e","[","0","]",",","rate=e","[","1","]",")","n","~","Normal","(","loc=0",",","scale=2",".",")","m","~","Normal","(","loc=n",",","scale=g",")","for","i","=","1",",","...",",","12",":","x","[","i","]","~","Bernoulli","(","logits=m",")","``","`","We","can","code","this","as",":","``","`","python","tfd","=","tfp.distributions","joint","=","tfd.JointDistributionNamedAutoBatched","(","dict","(","e=","tfd.Exponential","(","rate=","[","100",",","120","]",")",",","g=lambda","e",":","tfd.Gamma","(","concentration=e","[","0","]",",","rate=e","[","1","]",")",",","n=","tfd.Normal","(","loc=0",",","scale=2",".",")",",","m=lambda","n",",","g",":","tfd.Normal","(","loc=n",",","scale=g",")",",","x=lambda","m",":","tfd.Sample","(","tfd.Bernoulli","(","logits=m",")",",","12",")",",",")",")","``","`","Notice","the","1:1","correspondence","between","``","math","''","and","``","code","''",".","In","a","standard","`","JointDistributionNamed","`",",","we","would","have","wrapped","the","first","variable","as","`","e","=","tfd.Independent","(","tfd.Exponential","(","rate=","[","100",",","120","]",")",",","reinterpreted_batch_ndims=1",")","`","to","specify","that","`","log_prob","`","of","the","`","Exponential","`","should","be","a","scalar",",","summing","over","both","dimensions",".","We","would","also","have","had","to","extend","indices","as","`","tfd.Gamma","(","concentration=e","[","...",",","0","]",",","rate=e","[","...",",","1","]",")","`","to","account","for","possible","batch","dimensions",".","Both","of","these","behaviors","are","implicit","in","`","JointDistributionNamedAutoBatched","`","."]
docstr: Construct the `JointDistributionNamedAutoBatched` distribution.
  
  Args:
    model: Python `dict`, `collections.OrderedDict`, or `namedtuple` of
      distribution-making functions each with required args corresponding
      only to other keys.
    batch_ndims: `int` `Tensor` number of batch dimensions. The `batch_shape`s
      of all component distributions must be such that the prefixes of
      length `batch_ndims` broadcast to a consistent joint batch shape.
      Default value: `0`.
    use_vectorized_map: Python `bool`. Whether to use `tf.vectorized_map`
      to automatically vectorize evaluation of the model. This allows the
      model specification to focus on drawing a single sample, which is often
      simpler, but some ops may not be supported.
      Default value: `True`.
    validate_args: Python `bool`.  Whether to validate input with asserts.
      If `validate_args` is `False`, and the inputs are invalid,
      correct behavior is not guaranteed.
      Default value: `False`.
    experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
      summation to aggregate independent underlying log_prob values, which
      improves against the precision of a naive float32 sum. This can be
      noticeable in particular for large dimensions in float32. See CPU caveat
      on `tfp.math.reduce_kahan_sum`.
    name: The name for ops managed by the distribution.
      Default value: `None` (i.e., `JointDistributionNamed`).
docstr_tok: ["Construct","the","`","JointDistributionNamedAutoBatched","`","distribution",".","Args",":","model",":","Python","`","dict","`",",","`","collections.OrderedDict","`",",","or","`","namedtuple","`","of","distribution-making","functions","each","with","required","args","corresponding","only","to","other","keys",".","batch_ndims",":","`","int","`","`","Tensor","`","number","of","batch","dimensions",".","The","`","batch_shape","`","s","of","all","component","distributions","must","be","such","that","the","prefixes","of","length","`","batch_ndims","`","broadcast","to","a","consistent","joint","batch","shape",".","Default","value",":","`","0","`",".","use_vectorized_map",":","Python","`","bool","`",".","Whether","to","use","`","tf.vectorized_map","`","to","automatically","vectorize","evaluation","of","the","model",".","This","allows","the","model","specification","to","focus","on","drawing","a","single","sample",",","which","is","often","simpler",",","but","some","ops","may","not","be","supported",".","Default","value",":","`","True","`",".","validate_args",":","Python","`","bool","`",".","Whether","to","validate","input","with","asserts",".","If","`","validate_args","`","is","`","False","`",",","and","the","inputs","are","invalid",",","correct","behavior","is","not","guaranteed",".","Default","value",":","`","False","`",".","experimental_use_kahan_sum",":","Python","`","bool","`",".","When","`","True","`",",","we","use","Kahan","summation","to","aggregate","independent","underlying","log_prob","values",",","which","improves","against","the","precision","of","a","naive","float32","sum",".","This","can","be","noticeable","in","particular","for","large","dimensions","in","float32",".","See","CPU","caveat","on","`","tfp.math.reduce_kahan_sum","`",".","name",":","The","name","for","ops","managed","by","the","distribution",".","Default","value",":","`","None","`","(","i.e.",",","`","JointDistributionNamed","`",")","."]
code: def __init__(
      self,
      model,
      batch_ndims=0,
      use_vectorized_map=True,
      validate_args=False,
      experimental_use_kahan_sum=False,
      name=None,
  ):
      """Construct the `JointDistributionNamedAutoBatched` distribution.
  
      Args:
        model: Python `dict`, `collections.OrderedDict`, or `namedtuple` of
          distribution-making functions each with required args corresponding
          only to other keys.
        batch_ndims: `int` `Tensor` number of batch dimensions. The `batch_shape`s
          of all component distributions must be such that the prefixes of
          length `batch_ndims` broadcast to a consistent joint batch shape.
          Default value: `0`.
        use_vectorized_map: Python `bool`. Whether to use `tf.vectorized_map`
          to automatically vectorize evaluation of the model. This allows the
          model specification to focus on drawing a single sample, which is often
          simpler, but some ops may not be supported.
          Default value: `True`.
        validate_args: Python `bool`.  Whether to validate input with asserts.
          If `validate_args` is `False`, and the inputs are invalid,
          correct behavior is not guaranteed.
          Default value: `False`.
        experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
          summation to aggregate independent underlying log_prob values, which
          improves against the precision of a naive float32 sum. This can be
          noticeable in particular for large dimensions in float32. See CPU caveat
          on `tfp.math.reduce_kahan_sum`.
        name: The name for ops managed by the distribution.
          Default value: `None` (i.e., `JointDistributionNamed`).
      """
      parameters = dict(locals())
      super(_JointDistributionNamedAutoBatched, self).__init__(
          model,
          batch_ndims=batch_ndims,
          use_vectorized_map=use_vectorized_map,
          validate_args=validate_args,
          experimental_use_kahan_sum=experimental_use_kahan_sum,
          name=name or "JointDistributionNamedAutoBatched",
      )
      self._parameters = self._no_dependency(parameters)
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.internal import auto_composite_tensor
label_desc: Use image node_id 1 to create a new _JointDistributionNamedAutoBatched object from inherited base classes: joint_distribution_named with example: obj = _JointDistributionNamedAutoBatched(model, batch_ndims, use_vectorized_map, validate_args, experimental_use_kahan_sum, name)