docs: def build_modules(is_training, vocab_size):
    """Construct the modules used in the graph."""
  
    # Construct the custom getter which implements Bayes by Backprop.
    if is_training:
      estimator_mode = tf.constant(bbb.EstimatorModes.sample)
    else:
      estimator_mode = tf.constant(bbb.EstimatorModes.mean)
    lstm_bbb_custom_getter = bbb.bayes_by_backprop_getter(
        posterior_builder=lstm_posterior_builder,
        prior_builder=custom_scale_mixture_prior_builder,
        kl_builder=bbb.stochastic_kl_builder,
        sampling_mode_tensor=estimator_mode)
    non_lstm_bbb_custom_getter = bbb.bayes_by_backprop_getter(
        posterior_builder=non_lstm_posterior_builder,
        prior_builder=custom_scale_mixture_prior_builder,
        kl_builder=bbb.stochastic_kl_builder,
        sampling_mode_tensor=estimator_mode)
  
    embed_layer = snt.Embed(
        vocab_size=vocab_size,
        embed_dim=FLAGS.embedding_size,
        custom_getter=non_lstm_bbb_custom_getter,
        name="input_embedding")
  
    cores = []
    for i in range(FLAGS.n_layers):
      cores.append(
          snt.LSTM(FLAGS.hidden_size,
                   custom_getter=lstm_bbb_custom_getter,
                   forget_bias=0.0,
                   name="lstm_layer_{}".format(i)))
  
    rnn_core = snt.DeepRNN(
        cores,
        skip_connections=False,
        name="deep_lstm_core")
  
    # Do BBB on weights but not biases of output layer.
    output_linear = snt.Linear(
        vocab_size, custom_getter={"w": non_lstm_bbb_custom_getter})
    return embed_layer, rnn_core, output_linear