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: Adds reparameterization (pathwise) gradients to samples of the mixture.
  
  Implicit reparameterization gradients are
     dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi,
  where transform(x, phi) is distributional transform that removes all
  parameters from samples x.
  
  We implement them by replacing x with
    -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)]
  for the backward pass (gradient computation).
  The derivative of this quantity w.r.t. phi is then the implicit
  reparameterization gradient.
  Note that this replaces the gradients w.r.t. both the mixture
  distribution parameters and components distributions parameters.
  
  Limitations:
    1. Fundamental: components must be fully reparameterized.
    2. Distributional transform is currently only implemented for
      factorized components.
    3. Distributional transform currently only works for known rank of the
      batch tensor.
  
  Args:
    x: Sample of mixture distribution
    event_shape: The event shape of this distribution
  
  Returns:
    Tensor with same value as x, but with reparameterization gradients
docstr_tok: ["Adds","reparameterization","(","pathwise",")","gradients","to","samples","of","the","mixture",".","Implicit","reparameterization","gradients","are","dx\/dphi","=","-","(","d","transform","(","x",",","phi",")","\/","dx",")","^-1","*","d","transform","(","x",",","phi",")","\/","dphi",",","where","transform","(","x",",","phi",")","is","distributional","transform","that","removes","all","parameters","from","samples","x",".","We","implement","them","by","replacing","x","with","-stop_gradient","(","d","transform","(","x",",","phi",")","\/","dx",")","^-1","*","transform","(","x",",","phi",")","]","for","the","backward","pass","(","gradient","computation",")",".","The","derivative","of","this","quantity","w.r.t",".","phi","is","then","the","implicit","reparameterization","gradient",".","Note","that","this","replaces","the","gradients","w.r.t",".","both","the","mixture","distribution","parameters","and","components","distributions","parameters",".","Limitations",":","1",".","Fundamental",":","components","must","be","fully","reparameterized",".","2",".","Distributional","transform","is","currently","only","implemented","for","factorized","components",".","3",".","Distributional","transform","currently","only","works","for","known","rank","of","the","batch","tensor",".","Args",":","x",":","Sample","of","mixture","distribution","event_shape",":","The","event","shape","of","this","distribution","Returns",":","Tensor","with","same","value","as","x",",","but","with","reparameterization","gradients"]
code: def _reparameterize_sample(self, x, event_shape):
      """Adds reparameterization (pathwise) gradients to samples of the mixture.
  
      Implicit reparameterization gradients are
         dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi,
      where transform(x, phi) is distributional transform that removes all
      parameters from samples x.
  
      We implement them by replacing x with
        -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)]
      for the backward pass (gradient computation).
      The derivative of this quantity w.r.t. phi is then the implicit
      reparameterization gradient.
      Note that this replaces the gradients w.r.t. both the mixture
      distribution parameters and components distributions parameters.
  
      Limitations:
        1. Fundamental: components must be fully reparameterized.
        2. Distributional transform is currently only implemented for
          factorized components.
        3. Distributional transform currently only works for known rank of the
          batch tensor.
  
      Args:
        x: Sample of mixture distribution
        event_shape: The event shape of this distribution
  
      Returns:
        Tensor with same value as x, but with reparameterization gradients
      """
      # Remove the existing gradients of x wrt parameters of the components.
      x = tf.stop_gradient(x)
  
      event_size = ps.cast(ps.reduce_prod(event_shape), dtype=tf.int32)
      x_2d_shape = [-1, event_size]  # [S*prod(B), prod(E)]
  
      # Perform distributional transform of x in [S, B, E] shape,
      # but have Jacobian of size [S*prod(B), prod(E), prod(E)].
      def reshaped_distributional_transform(x_2d):
          return tf.reshape(
              self._distributional_transform(
                  tf.reshape(x_2d, ps.shape(x)), event_shape
              ),
              x_2d_shape,
          )
  
      # transform_2d: [S*prod(B), prod(E)]
      # jacobian: [S*prod(B), prod(E), prod(E)]
      x_2d = tf.reshape(x, x_2d_shape)
      transform_2d, jacobian = value_and_batch_jacobian(
          reshaped_distributional_transform, x_2d
      )
  
      # We only provide the first derivative; the second derivative computed by
      # autodiff would be incorrect, so we raise an error if it is requested.
      transform_2d = _prevent_2nd_derivative(transform_2d)
  
      # Avoid short-circuiting from `tf.stop_gradient`, which might otherwise hide
      # the 2nd derivative prevention.
      soft_stop_gradient = tf.custom_gradient(
          lambda v: (v, lambda dv: None)
      )
      # Compute [- stop_gradient(jacobian)^-1 * transform] by solving a linear
      # system. The Jacobian is lower triangular because the distributional
      # transform for i-th event dimension does not depend on the next
      # dimensions.
      surrogate_x_2d = -tf.linalg.triangular_solve(
          soft_stop_gradient(jacobian),
          transform_2d[..., tf.newaxis],
          lower=True,
      )  # [S*prod(B), prod(E), 1]
      surrogate_x = tf.reshape(surrogate_x_2d, ps.shape(x))
  
      # Replace gradients of x with gradients of surrogate_x, but keep the value.
      return x + (surrogate_x - tf.stop_gradient(surrogate_x))
  
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 20 for calling the _MixtureSameFamily obj's underlying member method code with example usage: obj._reparameterize_sample(x, event_shape) and returns: unknown, tf