class_docstr: Marginazlization mixin.
  
  This enables a joint distribution class to compute marginalized
  log probabilities.
docstr: Log probability density/mass function.
  
  Args:
    values: Structure of `Tensor`-values corresponding to samples yielded
      by model. There are also two special values that can be provided
      instead of `Tensor` samples.
      'marginalize': treat this random variable as a latent variable
      that will be marginalized out by summing over the support of
      the variable.
      'tabulate': treat this as a latent variable whose posterior probability
      distribution is to be computed. The final result is a tensor
      of log probabilities with each tabulated variable corresponding to
      one axis of the tensor in the order they appear in the list
      of `values`.
    name: Python `str` prepended to names of ops created by this function.
    method: Specifies method by which marginalization is carried out.
      'einsum': use `einsum`. For very small probabilities this may
      result in underflow causing log probabilities of `-inf` even when
      the result shoukd be finite.
      'logeinsumexp': performs an `einsum` designed to work in log space.
      Although it preserves precision better than the 'einsum' method it
      can be slow and use more memory.
    internal_type: because the `einsum` method can cause underflow this
      argument allows the user to specify the type in which the
      `einsum` is computed. For example `tf.float64`.
  
  Returns:
    Return log probability density/mass of `value` for this distribution.
  
  Notes:
    Currently only a single log probability can be computed, so lists or
    tensors containing multiple samples from the joint distribution are
    not supported.
    The number of latent (i.e. marginalized or tabulated) variables is
    limited to 52 in this version.
    The individual samples in `tfd.Sample` are mathematically independent
    but the marginalization algorithm used is unable to exploit this fact,
    meaning that computation time can easily grow exponentially with
    `sample_shape`.
docstr_tok: ["Log","probability","density\/mass","function",".","Args",":","values",":","Structure","of","`","Tensor","`","-values","corresponding","to","samples","yielded","by","model",".","There","are","also","two","special","values","that","can","be","provided","instead","of","`","Tensor","`","samples",".","'marginalize","'",":","treat","this","random","variable","as","a","latent","variable","that","will","be","marginalized","out","by","summing","over","the","support","of","the","variable",".","'tabulate","'",":","treat","this","as","a","latent","variable","whose","posterior","probability","distribution","is","to","be","computed",".","The","final","result","is","a","tensor","of","log","probabilities","with","each","tabulated","variable","corresponding","to","one","axis","of","the","tensor","in","the","order","they","appear","in","the","list","of","`","values","`",".","name",":","Python","`","str","`","prepended","to","names","of","ops","created","by","this","function",".","method",":","Specifies","method","by","which","marginalization","is","carried","out",".","'einsum","'",":","use","`","einsum","`",".","For","very","small","probabilities","this","may","result","in","underflow","causing","log","probabilities","of","`","-inf","`","even","when","the","result","shoukd","be","finite",".","'logeinsumexp","'",":","performs","an","`","einsum","`","designed","to","work","in","log","space",".","Although","it","preserves","precision","better","than","the","'einsum","'","method","it","can","be","slow","and","use","more","memory",".","internal_type",":","because","the","`","einsum","`","method","can","cause","underflow","this","argument","allows","the","user","to","specify","the","type","in","which","the","`","einsum","`","is","computed",".","For","example","`","tf.float64","`",".","Returns",":","Return","log","probability","density\/mass","of","`","value","`","for","this","distribution",".","Notes",":","Currently","only","a","single","log","probability","can","be","computed",",","so","lists","or","tensors","containing","multiple","samples","from","the","joint","distribution","are","not","supported",".","The","number","of","latent","(","i.e",".","marginalized","or","tabulated",")","variables","is","limited","to","52","in","this","version",".","The","individual","samples","in","`","tfd.Sample","`","are","mathematically","independent","but","the","marginalization","algorithm","used","is","unable","to","exploit","this","fact",",","meaning","that","computation","time","can","easily","grow","exponentially","with","`","sample_shape","`","."]
code: def marginalized_log_prob(
      self,
      values,
      name="marginalized_log_prob",
      method="logeinsumexp",
      internal_type=None,
  ):
      """Log probability density/mass function.
  
      Args:
        values: Structure of `Tensor`-values corresponding to samples yielded
          by model. There are also two special values that can be provided
          instead of `Tensor` samples.
          'marginalize': treat this random variable as a latent variable
          that will be marginalized out by summing over the support of
          the variable.
          'tabulate': treat this as a latent variable whose posterior probability
          distribution is to be computed. The final result is a tensor
          of log probabilities with each tabulated variable corresponding to
          one axis of the tensor in the order they appear in the list
          of `values`.
        name: Python `str` prepended to names of ops created by this function.
        method: Specifies method by which marginalization is carried out.
          'einsum': use `einsum`. For very small probabilities this may
          result in underflow causing log probabilities of `-inf` even when
          the result shoukd be finite.
          'logeinsumexp': performs an `einsum` designed to work in log space.
          Although it preserves precision better than the 'einsum' method it
          can be slow and use more memory.
        internal_type: because the `einsum` method can cause underflow this
          argument allows the user to specify the type in which the
          `einsum` is computed. For example `tf.float64`.
  
      Returns:
        Return log probability density/mass of `value` for this distribution.
  
      Notes:
        Currently only a single log probability can be computed, so lists or
        tensors containing multiple samples from the joint distribution are
        not supported.
        The number of latent (i.e. marginalized or tabulated) variables is
        limited to 52 in this version.
        The individual samples in `tfd.Sample` are mathematically independent
        but the marginalization algorithm used is unable to exploit this fact,
        meaning that computation time can easily grow exponentially with
        `sample_shape`.
      """
      new_values = []
      indices = []
      formula = []
      table_rhs = []
  
      # Number of independent variables created so far
      num_variables = 0
  
      with tf.name_scope(name):
          ds = self._call_execute_model(
              sample_and_trace_fn=jd_lib.trace_distributions_only,
              # Only used for tracing so can be fixed.
              seed=samplers.zeros_seed(),
          )
  
          # Both 'marginalize' and 'tabulate' indicate that
          # instead of using samples provided by the user, this method
          # instead provides a tensor containing the entire support of
          # the distribution. The tensors are constructed so that the
          # support for each distribution is indexed by a different axis.
          # At the end, the probabilities are computed using `tf.einsum`
          # and the unique axes means that each tabulated or marginalized
          # variable corresponds to one symbol used in the `einsum`.
          for value, dist in zip(values, ds):
              if value == "marginalize":
                  supp, rank = _support(dist)
                  r = ps.rank(supp)
                  num_new_variables = r - rank
                  # We can think of supp as being a tensor containing tensors,
                  # each of which is a draw from the distribution.
                  # `rank` is the rank of samples from the distribution.
                  # `num_new_variables` is the rank of the containing tensor.
                  # When we marginalize over a variable we want the sum
                  # over the containing tensor.
                  # So `num_new_variables` is the number of new indices needed.
                  # We use `expand_right` to ensure that each of these
                  # new indices is unique and independent of previous
                  # supports.
                  new_values.append(
                      _expand_right(
                          supp, n=num_variables, pos=-1 - rank
                      )
                  )
                  num_variables += num_new_variables
                  indices = np.arange(num_variables - 1, -1, -1)
                  # By *not* placing indices on the right of the '->' in
                  # the einsum below we ensure that this variable is
                  # used for reduction, in effect marginalizing over
                  # that variable.
                  formula.append(indices)
              elif value == "tabulate":
                  supp, rank = _support(dist)
                  r = ps.rank(supp)
                  if r is None:
                      raise ValueError(
                          "Need to be able to statically find rank of"
                          "support of random variable: {}".format(
                              str(dist)
                          )
                      )
                  num_new_variables = r - rank
                  new_values.append(
                      _expand_right(
                          supp, n=num_variables, pos=-1 - rank
                      )
                  )
                  num_variables += num_new_variables
                  indices = np.arange(num_variables - 1, -1, -1)
                  # The first elements of `indices` are the newly
                  # introduced variables.
                  new_indices = indices[:num_new_variables]
                  # By placing indices on the right of the '->' in
                  # the einsum below we ensure that this variable isn't
                  # reduced over and that instead it is tabulated
                  # in the resulting tensor.
                  table_rhs.extend(new_indices)
                  formula.append(indices)
              else:
                  new_values.append(
                      _expand_right(value, num_variables, -1)
                  )
                  indices = range(num_variables - 1, -1, -1)
                  formula.append(indices)
          formula = ["".join(map(_letter, f)) for f in formula]
          formula_rhs = "".join(map(_letter, table_rhs))
          formula = "{}->{}".format(",".join(formula), formula_rhs)
  
          # There is no `logsumexp_einsum`.
          # So use higher precision if user requests it.
          lpp = self.log_prob_parts(new_values)
          if internal_type:
              lpp = [tf.cast(x, dtype=internal_type) for x in lpp]
          if method == "logeinsumexp":
              return _squeezed_einsum(logeinsumexp, formula, *lpp)
          elif method == "einsum":
              return tf.math.log(
                  _squeezed_einsum(
                      tf.einsum, formula, *map(tf.exp, lpp)
                  )
              )
          else:
              raise ValueError(
                  "Unknown `marginalized_log_prob` method: '{}'.".format(
                      method
                  )
              )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import joint_distribution
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.marginalize.logeinsumexp import logeinsumexp
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 1 for calling the Marginalizable obj's underlying member method code with example usage: obj.marginalized_log_prob(values, name, method, internal_type) and returns: _squeezed_einsum, tf