def call(self, inputs):
    """Runs the model to generate a distribution `q(f | x_{1:T})`.

    This generates a list of batched MultivariateNormalDiag
    distributions using the output of the recurrent model at each
    timestep to parameterize each distribution.

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

    Returns:
      A batched MultivariateNormalDiag distribution with event shape
      [latent_size], batch shape [..., batch_size], and sample shape
      [sample_shape, ..., batch_size, latent_size].
    """
    # TODO(dusenberrymw): Remove these reshaping commands after b/113126249 is
    # fixed.
    collapsed_shape = tf.concat(([-1], tf.shape(input=inputs)[-2:]), axis=0)
    out = tf.reshape(inputs, collapsed_shape)  # (sample*batch_size, T, hidden)
    out = self.bilstm(out)  # (sample*batch_size, hidden)
    expanded_shape = tf.concat((tf.shape(input=inputs)[:-2], [-1]), axis=0)
    out = tf.reshape(out, expanded_shape)  # (sample, batch_size, hidden)
    out = self.output_layer(out)  # (sample, batch_size, 2*latent_size)
    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(f | x_{1:T})`.

    This generates a list of batched MultivariateNormalDiag
    distributions using the output of the recurrent model at each
    timestep to parameterize each distribution.

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

    Returns:
      A batched MultivariateNormalDiag distribution with event shape
      [latent_size], batch shape [..., batch_size], and sample shape
      [sample_shape, ..., batch_size, latent_size].