Predict a sequence of bits (a latent) with LSTM, both training and infer.

  Given a tensor on which the predictions are based (prediction_source), we use
  a single-layer LSTM with state of size state_size to predict total_num_bits,
  which we predict in groups of size bits_at_once. During training, we use
  target_bits as input to the LSTM (teacher forcing) and return the target_bits
  together with the prediction loss. During inference, we sample with the given
  temperature and return the predicted sequence and loss 0.

  Args:
    prediction_source: a Tensor of shape [batch_size, ...] used to create
      the initial state and the first input to the LSTM.
    state_size: python integer, the size of the LSTM state.
    total_num_bits: python integer, how many bits in total to predict.
    target_bits: a tensor of shape [batch_size, total_num_bits] used during
      training as the target to predict; each element should be -1 or 1.
    extra_inputs: a Tensor [batch_size, total_num_bits // bits_at_once, d]
      of additional inputs, passed as additional LSTM inputs.
    bits_at_once: pytho integer, how many bits to predict at once.
    temperature: python float, temperature used for sampling during inference.
    dropout: float, the amount of dropout to aply during training (0.1 default).

  Returns:
    a pair (bits, loss) with the predicted bit sequence, which is a Tensor of
    shape [batch_size, total_num_bits] with elements either -1 or 1, and a loss
    used to train the predictions against the provided target_bits.

def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
                           target_bits=None, extra_inputs=None,
                           bits_at_once=8, temperature=1.0, dropout=0.1):
  """Predict a sequence of bits (a latent) with LSTM, both training and infer.

  Given a tensor on which the predictions are based (prediction_source), we use
  a single-layer LSTM with state of size state_size to predict total_num_bits,
  which we predict in groups of size bits_at_once. During training, we use
  target_bits as input to the LSTM (teacher forcing) and return the target_bits
  together with the prediction loss. During inference, we sample with the given
  temperature and return the predicted sequence and loss 0.

  Args:
    prediction_source: a Tensor of shape [batch_size, ...] used to create
      the initial state and the first input to the LSTM.
    state_size: python integer, the size of the LSTM state.
    total_num_bits: python integer, how many bits in total to predict.
    target_bits: a tensor of shape [batch_size, total_num_bits] used during
      training as the target to predict; each element should be -1 or 1.
    extra_inputs: a Tensor [batch_size, total_num_bits // bits_at_once, d]
      of additional inputs, passed as additional LSTM inputs.
    bits_at_once: pytho integer, how many bits to predict at once.
    temperature: python float, temperature used for sampling during inference.
    dropout: float, the amount of dropout to aply during training (0.1 default).

  Returns:
    a pair (bits, loss) with the predicted bit sequence, which is a Tensor of
    shape [batch_size, total_num_bits] with elements either -1 or 1, and a loss
    used to train the predictions against the provided target_bits.
  """

  with tf.variable_scope("predict_bits_with_lstm"):
    # Layers and cell state creation.
    lstm_cell = tf.nn.rnn_cell.LSTMCell(state_size)
    discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
    discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
    batch_size = common_layers.shape_list(prediction_source)[0]
    layer_pred = tf.layers.flatten(prediction_source)
    first_lstm_input = tf.layers.dense(layer_pred, state_size, name="istate")
    c_state = tf.layers.dense(layer_pred, state_size, name="cstate")
    m_state = tf.layers.dense(layer_pred, state_size, name="mstate")
    state = (c_state, m_state)

    # Prediction mode if no targets are given.
    if target_bits is None:
      outputs = []
      lstm_input = first_lstm_input
      for i in range(total_num_bits // bits_at_once):
        if extra_inputs is not None:
          lstm_input = tf.concat([lstm_input, extra_inputs[:, i, :]], axis=1)
        output, state = lstm_cell(lstm_input, state)
        discrete_logits = discrete_predict(output)
        discrete_samples = common_layers.sample_with_temperature(
            discrete_logits, temperature)
        outputs.append(tf.expand_dims(discrete_samples, axis=1))
        lstm_input = discrete_embed(tf.one_hot(discrete_samples, 256))
      outputs = tf.concat(outputs, axis=1)
      outputs = int_to_bit(outputs, bits_at_once)
      outputs = tf.reshape(outputs, [batch_size, total_num_bits])
      return 2 * outputs - 1, 0.0

    # Training mode, calculating loss.
    assert total_num_bits % bits_at_once == 0
    target_bits = tf.reshape(tf.maximum(tf.stop_gradient(target_bits), 0), [
        batch_size, total_num_bits // bits_at_once, bits_at_once])
    target_ints = bit_to_int(target_bits, bits_at_once)
    tf.summary.histogram("target_integers", tf.reshape(target_ints, [-1]))
    target_hot = tf.one_hot(target_ints, 2**bits_at_once, axis=-1)
    target_embedded = discrete_embed(target_hot)
    target_embedded = tf.nn.dropout(target_embedded, 1.0 - dropout)
    teacher_input = tf.concat(
        [tf.expand_dims(first_lstm_input, axis=1), target_embedded], axis=1)
    outputs = []
    for i in range(total_num_bits // bits_at_once):
      lstm_input = teacher_input[:, i, :]
      if extra_inputs is not None:
        lstm_input = tf.concat([lstm_input, extra_inputs[:, i, :]], axis=1)
      output, state = lstm_cell(lstm_input, state)
      outputs.append(tf.expand_dims(output, axis=1))
    outputs = tf.concat(outputs, axis=1)
    outputs = tf.nn.dropout(outputs, 1.0 - dropout)
    d_int_pred = discrete_predict(outputs)
    pred_loss = tf.losses.sparse_softmax_cross_entropy(
        logits=d_int_pred, labels=target_ints)
    pred_loss = tf.reduce_mean(pred_loss)
    return d_int_pred, pred_loss