import tensorflow as tf


    Args:
      prev_rnn_state: The state of the RNN from the previous timestep.
      prev_latent_encoded: Float Tensor of shape
        [batch_size, encoded_latent_size], the previous latent state z_{t-1}
        run through latent_encoder.
      inputs: A Tensor of shape [batch_size, data_size], the current inputs to
        the model. Most often this is x_{t-1}, the previous token in the
        observation sequence.
    Returns:
      rnn_out: The output of the RNN.
      rnn_state: The new state of the RNN.
    """
    inputs_encoded = self.data_encoder(tf.to_float(inputs))
    rnn_inputs = tf.concat([inputs_encoded, prev_latent_encoded], axis=1)
    rnn_out, rnn_state = self.rnn_cell(rnn_inputs, prev_rnn_state)
    return rnn_out, rnn_state

  def transition(self, rnn_out):
    """Computes the transition distribution p(z_t|h_t).

    Note that p(z_t | h_t) = p(z_t| z_{1:t-1}, x_{1:t-1})

    Args:
      rnn_out: The output of the rnn for the current timestep.
    Returns:
