class_docstr: Mixture (same-family) distribution.
  
  The `MixtureSameFamily` distribution implements a (batch of) mixture
  distribution where all components are from different parameterizations of the
  same distribution type. It is parameterized by a `Categorical` 'selecting
  distribution' (over `k` components) and a components distribution, i.e., a
  `Distribution` with a rightmost batch shape (equal to `[k]`) which indexes
  each (batch of) component.
  
  #### Examples
  
  ```python
  tfd = tfp.distributions
  
  ### Create a mixture of two scalar Gaussians:
  
  gm = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(
          probs=[0.3, 0.7]),
      components_distribution=tfd.Normal(
        loc=[-1., 1],       # One for each component.
        scale=[0.1, 0.5]))  # And same here.
  
  gm.mean()
  # ==> 0.4
  
  gm.variance()
  # ==> 1.018
  
  # Plot PDF.
  x = np.linspace(-2., 3., int(1e4), dtype=np.float32)
  import matplotlib.pyplot as plt
  plt.plot(x, gm.prob(x));
  
  ### Create a mixture of three Bivariate Gaussians:
  
  gm = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(
          probs=[0.2, 0.4, 0.4]),
      components_distribution=tfd.MultivariateNormalDiag(
          loc=[[-1., 1],  # component 1
               [1, -1],  # component 2
               [1, 1]],  # component 3
          scale_diag=tf.tile([[.3], [.6], [.7]], [1, 2])))
  
  gm.components_distribution.batch_shape
  # ==> (3,)
  
  gm.components_distribution.event_shape
  # ==> (2,)
  
  gm.mean()
  # ==> array([ 0.6, 0.2], dtype=float32)
  
  gm.covariance()
  # ==> array([[ 0.998    , -0.32     ],
  #            [-0.32     ,  1.3180001]], dtype=float32)
  
  # Plot PDF contours.
  def meshgrid(x):
    y = x
    [gx, gy] = np.meshgrid(x, y, indexing='ij')
    gx, gy = np.float32(gx), np.float32(gy)
    grid = np.concatenate([gx.ravel()[None, :], gy.ravel()[None, :]], axis=0)
    return grid.T.reshape(x.size, y.size, 2)
  grid = meshgrid(np.linspace(-2, 2, 100, dtype=np.float32))
  plt.contour(grid[..., 0], grid[..., 1], gm.prob(grid));
  ```
  
  Note that this distribution is *not* a joint distribution over categorical
  and continuous values, but rather a mixture of continuous distributions
  proportioned by the given categorical distribution. If you want a joint
  distribution, you might write it as:
  
  ```python
  @tfd.JointDistributionCoroutineAutoBatched
  def model():
    mus = tf.constant([[-1., 1], # component 1
                       [1, -1],  # component 2
                       [1, 1]])  # component 3
    scales = tf.constant([.3, .6, .7])
    idx = yield tfd.Categorical(probs=[.2, .4, .4], name='idx')
    val = yield tfd.MultivariateNormalDiag(
        loc=mus[idx], scale_diag=tf.ones(2) * scales[idx], name='val')
  
  model.sample()
  # ==> StructTuple(
  #       idx=2,
  #       val=array([1.0582672, 1.3583777], dtype=float32)
  #     )
  ```
class_docstr_tok: ["Mixture","(","same-family",")","distribution",".","The","`","MixtureSameFamily","`","distribution","implements","a","(","batch","of",")","mixture","distribution","where","all","components","are","from","different","parameterizations","of","the","same","distribution","type",".","It","is","parameterized","by","a","`","Categorical","`","'selecting","distribution","'","(","over","`","k","`","components",")","and","a","components","distribution",",","i.e.",",","a","`","Distribution","`","with","a","rightmost","batch","shape","(","equal","to","`","[","k","]","`",")","which","indexes","each","(","batch","of",")","component",".","#","#","#","#","Examples","``","`","python","tfd","=","tfp.distributions","#","#","#","Create","a","mixture","of","two","scalar","Gaussians",":","gm","=","tfd.MixtureSameFamily","(","mixture_distribution=tfd.Categorical","(","probs=","[","0.3",",","0.7","]",")",",","components_distribution=tfd.Normal","(","loc=","[","-1.",",","1","]",",","#","One","for","each","component",".","scale=","[","0.1",",","0.5","]",")",")","#","And","same","here",".","gm.mean","(",")","#","==",">","0.4","gm.variance","(",")","#","==",">","1.018","#","Plot","PDF",".","x","=","np.linspace","(","-2.",",","3.",",","int","(","1e4",")",",","dtype=np.float32",")","import","matplotlib.pyplot","as","plt","plt.plot","(","x",",","gm.prob","(","x",")",")",";","#","#","#","Create","a","mixture","of","three","Bivariate","Gaussians",":","gm","=","tfd.MixtureSameFamily","(","mixture_distribution=tfd.Categorical","(","probs=","[","0.2",",","0.4",",","0.4","]",")",",","components_distribution=tfd.MultivariateNormalDiag","(","loc=","[","[","-1.",",","1","]",",","#","component","1","[","1",",","-1","]",",","#","component","2","[","1",",","1","]","]",",","#","component","3","scale_diag=tf.tile","(","[","[",".3","]",",","[",".6","]",",","[",".7","]","]",",","[","1",",","2","]",")",")",")","gm.components_distribution.batch_shape","#","==",">","(","3",",",")","gm.components_distribution.event_shape","#","==",">","(","2",",",")","gm.mean","(",")","#","==",">","array","(","[","0.6",",","0.2","]",",","dtype=float32",")","gm.covariance","(",")","#","==",">","array","(","[","[","0.998",",","-0.32","]",",","#","[","-0.32",",","1.3180001","]","]",",","dtype=float32",")","#","Plot","PDF","contours",".","def","meshgrid","(","x",")",":","y","=","x","[","gx",",","gy","]","=","np.meshgrid","(","x",",","y",",","indexing='ij","'",")","gx",",","gy","=","np.float32","(","gx",")",",","np.float32","(","gy",")","grid","=","np.concatenate","(","[","gx.ravel","(",")","[","None",",",":","]",",","gy.ravel","(",")","[","None",",",":","]","]",",","axis=0",")","return","grid.T.reshape","(","x.size",",","y.size",",","2",")","grid","=","meshgrid","(","np.linspace","(","-2",",","2",",","100",",","dtype=np.float32",")",")","plt.contour","(","grid","[","...",",","0","]",",","grid","[","...",",","1","]",",","gm.prob","(","grid",")",")",";","``","`","Note","that","this","distribution","is","*","not","*","a","joint","distribution","over","categorical","and","continuous","values",",","but","rather","a","mixture","of","continuous","distributions","proportioned","by","the","given","categorical","distribution",".","If","you","want","a","joint","distribution",",","you","might","write","it","as",":","``","`","python","@","tfd.JointDistributionCoroutineAutoBatched","def","model","(",")",":","mus","=","tf.constant","(","[","[","-1.",",","1","]",",","#","component","1","[","1",",","-1","]",",","#","component","2","[","1",",","1","]","]",")","#","component","3","scales","=","tf.constant","(","[",".3",",",".6",",",".7","]",")","idx","=","yield","tfd.Categorical","(","probs=","[",".2",",",".4",",",".4","]",",","name='idx","'",")","val","=","yield","tfd.MultivariateNormalDiag","(","loc=mus","[","idx","]",",","scale_diag=tf.ones","(","2",")","*","scales","[","idx","]",",","name='val","'",")","model.sample","(",")","#","==",">","StructTuple","(","#","idx=2",",","#","val=array","(","[","1.0582672",",","1.3583777","]",",","dtype=float32",")","#",")","``","`"]
docstr: Construct a `MixtureSameFamily` distribution.
  
  Args:
    mixture_distribution: `tfd.Categorical`-like instance.
      Manages the probability of selecting components. The number of
      categories must match the rightmost batch dimension of the
      `components_distribution`. Must have `batch_shape` broadcastable
      with `components_distribution.batch_shape[:-1]`.
    components_distribution: `tfd.Distribution`-like instance.
      The right-most batch dimension indexes the mixture components.
    reparameterize: Python `bool`, default `False`. Whether to reparameterize
      samples of the distribution using implicit reparameterization gradients
      [(Figurnov et al., 2018)][1]. The gradients for the mixture logits are
      equivalent to the ones described by [(Graves, 2016)][2]. The gradients
      for the components parameters are also computed using implicit
      reparameterization (as opposed to ancestral sampling), meaning that
      all components are updated every step.
      Only works when:
        (1) components_distribution is fully reparameterized;
        (2) components_distribution is either a scalar distribution or
        fully factorized (tfd.Independent applied to a scalar distribution);
        (3) batch shape has a known rank.
      Experimental, may be slow and produce infs/NaNs.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
      (e.g., mean, mode, variance) use the value '`NaN`' to indicate the
      result is undefined. When `False`, an exception is raised if one or
      more of the statistic's batch members are undefined.
    name: Python `str` name prefixed to Ops created by this class.
  
  Raises:
    ValueError: `if not dtype_util.is_integer(mixture_distribution.dtype)`.
    ValueError: if mixture_distribution does not have scalar `event_shape`.
    ValueError: if `mixture_distribution.batch_shape` and
      `components_distribution.batch_shape[:-1]` are both fully defined and
      the former is neither scalar nor equal to the latter.
    ValueError: if `mixture_distribution` categories does not equal
      `components_distribution` rightmost batch shape.
  
  #### References
  
  [1]: Michael Figurnov, Shakir Mohamed and Andriy Mnih. Implicit
       reparameterization gradients. In _Neural Information Processing
       Systems_, 2018. https://arxiv.org/abs/1805.08498
  
  [2]: Alex Graves. Stochastic Backpropagation through Mixture Density
       Distributions. _arXiv_, 2016. https://arxiv.org/abs/1607.05690
docstr_tok: ["Construct","a","`","MixtureSameFamily","`","distribution",".","Args",":","mixture_distribution",":","`","tfd.Categorical","`","-like","instance",".","Manages","the","probability","of","selecting","components",".","The","number","of","categories","must","match","the","rightmost","batch","dimension","of","the","`","components_distribution","`",".","Must","have","`","batch_shape","`","broadcastable","with","`","components_distribution.batch_shape","[",":","-1","]","`",".","components_distribution",":","`","tfd.Distribution","`","-like","instance",".","The","right-most","batch","dimension","indexes","the","mixture","components",".","reparameterize",":","Python","`","bool","`",",","default","`","False","`",".","Whether","to","reparameterize","samples","of","the","distribution","using","implicit","reparameterization","gradients","[","(","Figurnov","et","al.",",","2018",")","]","[","1","]",".","The","gradients","for","the","mixture","logits","are","equivalent","to","the","ones","described","by","[","(","Graves",",","2016",")","]","[","2","]",".","The","gradients","for","the","components","parameters","are","also","computed","using","implicit","reparameterization","(","as","opposed","to","ancestral","sampling",")",",","meaning","that","all","components","are","updated","every","step",".","Only","works","when",":","(","1",")","components_distribution","is","fully","reparameterized",";","(","2",")","components_distribution","is","either","a","scalar","distribution","or","fully","factorized","(","tfd.Independent","applied","to","a","scalar","distribution",")",";","(","3",")","batch","shape","has","a","known","rank",".","Experimental",",","may","be","slow","and","produce","infs\/NaNs",".","validate_args",":","Python","`","bool","`",",","default","`","False","`",".","When","`","True","`","distribution","parameters","are","checked","for","validity","despite","possibly","degrading","runtime","performance",".","When","`","False","`","invalid","inputs","may","silently","render","incorrect","outputs",".","allow_nan_stats",":","Python","`","bool","`",",","default","`","True","`",".","When","`","True","`",",","statistics","(","e.g.",",","mean",",","mode",",","variance",")","use","the","value","'","`","NaN","`","'","to","indicate","the","result","is","undefined",".","When","`","False","`",",","an","exception","is","raised","if","one","or","more","of","the","statistic","'s","batch","members","are","undefined",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","Raises",":","ValueError",":","`","if","not","dtype_util.is_integer","(","mixture_distribution.dtype",")","`",".","ValueError",":","if","mixture_distribution","does","not","have","scalar","`","event_shape","`",".","ValueError",":","if","`","mixture_distribution.batch_shape","`","and","`","components_distribution.batch_shape","[",":","-1","]","`","are","both","fully","defined","and","the","former","is","neither","scalar","nor","equal","to","the","latter",".","ValueError",":","if","`","mixture_distribution","`","categories","does","not","equal","`","components_distribution","`","rightmost","batch","shape",".","#","#","#","#","References","[","1","]",":","Michael","Figurnov",",","Shakir","Mohamed","and","Andriy","Mnih",".","Implicit","reparameterization","gradients",".","In","_Neural","Information","Processing","Systems_",",","2018.","https",":","\/\/arxiv.org\/abs\/1805.08498","[","2","]",":","Alex","Graves",".","Stochastic","Backpropagation","through","Mixture","Density","Distributions",".","_arXiv_",",","2016.","https",":","\/\/arxiv.org\/abs\/1607.05690"]
code: def __init__(
      self,
      mixture_distribution,
      components_distribution,
      reparameterize=False,
      validate_args=False,
      allow_nan_stats=True,
      name="MixtureSameFamily",
  ):
      """Construct a `MixtureSameFamily` distribution.
  
      Args:
        mixture_distribution: `tfd.Categorical`-like instance.
          Manages the probability of selecting components. The number of
          categories must match the rightmost batch dimension of the
          `components_distribution`. Must have `batch_shape` broadcastable
          with `components_distribution.batch_shape[:-1]`.
        components_distribution: `tfd.Distribution`-like instance.
          The right-most batch dimension indexes the mixture components.
        reparameterize: Python `bool`, default `False`. Whether to reparameterize
          samples of the distribution using implicit reparameterization gradients
          [(Figurnov et al., 2018)][1]. The gradients for the mixture logits are
          equivalent to the ones described by [(Graves, 2016)][2]. The gradients
          for the components parameters are also computed using implicit
          reparameterization (as opposed to ancestral sampling), meaning that
          all components are updated every step.
          Only works when:
            (1) components_distribution is fully reparameterized;
            (2) components_distribution is either a scalar distribution or
            fully factorized (tfd.Independent applied to a scalar distribution);
            (3) batch shape has a known rank.
          Experimental, may be slow and produce infs/NaNs.
        validate_args: Python `bool`, default `False`. When `True` distribution
          parameters are checked for validity despite possibly degrading runtime
          performance. When `False` invalid inputs may silently render incorrect
          outputs.
        allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
          (e.g., mean, mode, variance) use the value '`NaN`' to indicate the
          result is undefined. When `False`, an exception is raised if one or
          more of the statistic's batch members are undefined.
        name: Python `str` name prefixed to Ops created by this class.
  
      Raises:
        ValueError: `if not dtype_util.is_integer(mixture_distribution.dtype)`.
        ValueError: if mixture_distribution does not have scalar `event_shape`.
        ValueError: if `mixture_distribution.batch_shape` and
          `components_distribution.batch_shape[:-1]` are both fully defined and
          the former is neither scalar nor equal to the latter.
        ValueError: if `mixture_distribution` categories does not equal
          `components_distribution` rightmost batch shape.
  
      #### References
  
      [1]: Michael Figurnov, Shakir Mohamed and Andriy Mnih. Implicit
           reparameterization gradients. In _Neural Information Processing
           Systems_, 2018. https://arxiv.org/abs/1805.08498
  
      [2]: Alex Graves. Stochastic Backpropagation through Mixture Density
           Distributions. _arXiv_, 2016. https://arxiv.org/abs/1607.05690
      """
      parameters = dict(locals())
      with tf.name_scope(name) as name:
          self._mixture_distribution = mixture_distribution
          self._components_distribution = components_distribution
  
          self._reparameterize = reparameterize
          if reparameterize:
              # Note: tfd.Independent passes through the reparameterization type hence
              # we do not need separate logic for Independent.
              if (
                  self._components_distribution.reparameterization_type
                  != reparameterization.FULLY_REPARAMETERIZED
              ):
                  raise ValueError(
                      "Cannot reparameterize a mixture of "
                      "non-reparameterized components."
                  )
              reparameterization_type = (
                  reparameterization.FULLY_REPARAMETERIZED
              )
          else:
              reparameterization_type = (
                  reparameterization.NOT_REPARAMETERIZED
              )
  
          super(_MixtureSameFamily, self).__init__(
              dtype=self._components_distribution.dtype,
              reparameterization_type=reparameterization_type,
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              parameters=parameters,
              name=name,
          )
  
imports: import warnings
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.gradient import value_and_batch_jacobian
label_desc: Use image node_id 1 to create a new _MixtureSameFamily object from inherited base classes: distribution with example: obj = _MixtureSameFamily(mixture_distribution, components_distribution, reparameterize, validate_args, allow_nan_stats, name)