tensorflow_probability/python/distributions/mixture_same_family.py
def _reparameterize_sample(self, x):
    """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.

    Arguments:
      x: Sample of mixture 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)

    x_2d_shape = [-1, self._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, tf.shape(input=x))),
          x_2d_shape)

    # transform_2d: [S*prod(B), prod(E)]
    # jacobian: [S*prod(B), prod(E), prod(E)]
    transform_2d, jacobian = _value_and_batch_jacobian(
        reshaped_distributional_transform, tf.reshape(x, x_2d_shape))

    # 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)

    # 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(
        tf.stop_gradient(jacobian), tf.expand_dims(transform_2d, axis=-1),
        lower=True)  # [S*prod(B), prod(E), 1]
    surrogate_x = tf.reshape(surrogate_x_2d, tf.shape(input=x))

    # Replace gradients of x with gradients of surrogate_x, but keep the value.
    return x + (surrogate_x - tf.stop_gradient(surrogate_x))
def _reparameterize_sample(self, x):
    """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.

    Arguments:
      x: Sample of mixture 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)

    x_2d_shape = [-1, self._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, tf.shape(input=x))),
          x_2d_shape)

    # transform_2d: [S*prod(B), prod(E)]
    # jacobian: [S*prod(B), prod(E), prod(E)]
    transform_2d, jacobian = _value_and_batch_jacobian(
        reshaped_distributional_transform, tf.reshape(x, x_2d_shape))

    # 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)

    # 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(
        tf.stop_gradient(jacobian), tf.expand_dims(transform_2d, axis=-1),
        lower=True)  # [S*prod(B), prod(E), 1]
    surrogate_x = tf.reshape(surrogate_x_2d, tf.shape(input=x))

    # Replace gradients of x with gradients of surrogate_x, but keep the value.
    return x + (surrogate_x - tf.stop_gradient(surrogate_x))
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.

    Arguments:
      x: Sample of mixture distribution

    Returns:
      Tensor with same value as x, but with reparameterization gradients