Construct the modules used in the graph.

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