tensorflow_probability/examples/disentangled_vae.py
def call(self, inputs):
    """Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].

    Returns:
      A batch of MultivariateNormalDiag distributions with event shape
      [latent_size], batch shape [..., batch_size, timesteps], and
      sample shape [sample_shape, ..., batch_size, timesteps,
      latent_size].
    """
    out = self.dense(inputs)  # (..., batch, time, hidden)
    out = self.output_layer(out)  # (..., batch, time, 2*latent)
    loc = out[..., :self.latent_size]
    scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
def call(self, inputs):
    """Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].

    Returns:
      A batch of MultivariateNormalDiag distributions with event shape
      [latent_size], batch shape [..., batch_size, timesteps], and
      sample shape [sample_shape, ..., batch_size, timesteps,
      latent_size].
    """
    out = self.dense(inputs)  # (..., batch, time, hidden)
    out = self.output_layer(out)  # (..., batch, time, 2*latent)
    loc = out[..., :self.latent_size]
    scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

    Args:
      inputs: A batch of intermediate representations of image frames
        across all timesteps, of shape [..., batch_size, timesteps,
        hidden_size].

    Returns:
      A batch of MultivariateNormalDiag distributions with event shape
      [latent_size], batch shape [..., batch_size, timesteps], and
      sample shape [sample_shape, ..., batch_size, timesteps,
      latent_size].