code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 14 for calling a global function with example usage: predict_bits_with_lstm(prediction_source, state_size, total_num_bits, target_bits, extra_inputs, bits_at_once, temperature, dropout) and returns: d_int_pred,  pred_loss, unknown,  int