docs: 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