tensorflow_probability/examples/disentangled_vae.py
def sample_dynamic_posterior(self, inputs, samples, static_sample=None):
    """Sample the static latent posterior.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].
      samples: Number of samples to draw from the latent distribution.
      static_sample: A tensor sample of the static latent variable `f`
        of shape [..., batch_size, latent_size]. Only used
        for the full dynamic posterior formulation.

    Returns:
      A tuple of a sample tensor of shape [samples, batch_size, length
      latent_size], and a MultivariateNormalDiag distribution from which
      the tensor was sampled, with event shape [latent_size], and batch
      shape [broadcasted_shape, batch_size, length], where
      `broadcasted_shape` is the broadcasted sampled shape between the
      inputs and static sample.

    Raises:
      ValueError: If the "full" latent posterior formulation is being
        used, yet a static latent sample was not provided.
    """
    if self.latent_posterior == "factorized":
      dist = self.dynamic_encoder(inputs)
      samples = dist.sample(samples)  # (s, N, T, lat)
    else:  # full
      if static_sample is None:
        raise ValueError(
            "The full dynamic posterior requires a static latent sample")
      dist = self.dynamic_encoder((inputs, static_sample))
      samples = dist.sample()  # (samples, N, latent)
    return samples, dist
def sample_dynamic_posterior(self, inputs, samples, static_sample=None):
    """Sample the static latent posterior.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].
      samples: Number of samples to draw from the latent distribution.
      static_sample: A tensor sample of the static latent variable `f`
        of shape [..., batch_size, latent_size]. Only used
        for the full dynamic posterior formulation.

    Returns:
      A tuple of a sample tensor of shape [samples, batch_size, length
      latent_size], and a MultivariateNormalDiag distribution from which
      the tensor was sampled, with event shape [latent_size], and batch
      shape [broadcasted_shape, batch_size, length], where
      `broadcasted_shape` is the broadcasted sampled shape between the
      inputs and static sample.

    Raises:
      ValueError: If the "full" latent posterior formulation is being
        used, yet a static latent sample was not provided.
    """
    if self.latent_posterior == "factorized":
      dist = self.dynamic_encoder(inputs)
      samples = dist.sample(samples)  # (s, N, T, lat)
    else:  # full
      if static_sample is None:
        raise ValueError(
            "The full dynamic posterior requires a static latent sample")
      dist = self.dynamic_encoder((inputs, static_sample))
      samples = dist.sample()  # (samples, N, latent)
    return samples, dist
Sample the static latent posterior.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].
      samples: Number of samples to draw from the latent distribution.
      static_sample: A tensor sample of the static latent variable `f`
        of shape [..., batch_size, latent_size]. Only used
        for the full dynamic posterior formulation.

    Returns:
      A tuple of a sample tensor of shape [samples, batch_size, length
      latent_size], and a MultivariateNormalDiag distribution from which
      the tensor was sampled, with event shape [latent_size], and batch
      shape [broadcasted_shape, batch_size, length], where
      `broadcasted_shape` is the broadcasted sampled shape between the
      inputs and static sample.

    Raises:
      ValueError: If the "full" latent posterior formulation is being
        used, yet a static latent sample was not provided.