This is the core model logic.

  Unrolls a Bayesian RNN over the given sequence.

  Args:
    data_ops: A `sequence_data.SequenceDataOps` namedtuple.
    embed_layer: A `snt.Embed` instance.
    rnn_core: A `snt.RNNCore` instance.
    output_linear: A `snt.Linear` instance.
    name_prefix: A string to use to prefix local variable names.

  Returns:
    A 3D time-major tensor representing the model's logits for a sequence of
    predictions. Shape `[time_steps, batch_size, vocab_size]`.

def build_logits(data_ops, embed_layer, rnn_core, output_linear, name_prefix):
  """This is the core model logic.

  Unrolls a Bayesian RNN over the given sequence.

  Args:
    data_ops: A `sequence_data.SequenceDataOps` namedtuple.
    embed_layer: A `snt.Embed` instance.
    rnn_core: A `snt.RNNCore` instance.
    output_linear: A `snt.Linear` instance.
    name_prefix: A string to use to prefix local variable names.

  Returns:
    A 3D time-major tensor representing the model's logits for a sequence of
    predictions. Shape `[time_steps, batch_size, vocab_size]`.
  """
  # Embed the input index sequence.
  embedded_input_seq = snt.BatchApply(
      embed_layer, name="input_embed_seq")(data_ops.sparse_obs)

  # Construct variables for holding the RNN state.
  initial_rnn_state = nest.map_structure(
      lambda t: tf.get_local_variable(  # pylint: disable long lambda warning
          "{}/rnn_state/{}".format(name_prefix, t.op.name), initializer=t),
      rnn_core.initial_state(FLAGS.batch_size))
  assign_zero_rnn_state = nest.map_structure(
      lambda x: x.assign(tf.zeros_like(x)), initial_rnn_state)
  assign_zero_rnn_state = tf.group(*nest.flatten(assign_zero_rnn_state))

  # Unroll the RNN core over the sequence.
  rnn_output_seq, rnn_final_state = tf.nn.dynamic_rnn(
      cell=rnn_core,
      inputs=embedded_input_seq,
      initial_state=initial_rnn_state,
      time_major=True)

  # Persist the RNN state for the next unroll.
  update_rnn_state = nest.map_structure(
      tf.assign, initial_rnn_state, rnn_final_state)
  with tf.control_dependencies(nest.flatten(update_rnn_state)):
    rnn_output_seq = tf.identity(rnn_output_seq, name="rnn_output_seq")
  output_logits = snt.BatchApply(
      output_linear, name="output_embed_seq")(rnn_output_seq)
  return output_logits, assign_zero_rnn_state