code: def train_step(self, input_data):
      # Note(Hao): enforce the lookup and lookup init to be replicated
      # with ops.device('/cpu:0'):
      text_init = tf.lookup.TextFileInitializer(
          os.path.join(FLAGS.datadir, "1b_word_vocab.txt"),
          tf.string,
          0,
          tf.int64,
          -1,
          delimiter=" ",
      )
      self.vocab = tf.lookup.StaticVocabularyTable(
          text_init, FLAGS.oov_bucket_size
      )
      xy = tf.dtypes.cast(self.vocab.lookup(input_data), dtype=tf.int32)
      x = xy[:, 0, :]
      y = xy[:, 1, :]
      w = np.ones([FLAGS.batch_size, FLAGS.num_steps])
  
      # FIXME: cast a PartitionVariabvles to a list of Variables
      # so that we can derive its gradients
      # otherwise will throw an error
      emb_vars = [self.emb]
      lstm_vars = [self.W, self.B, self.W_P]
      softmax_vars = [self.softmax_w, self.softmax_b]
      all_vars = emb_vars + lstm_vars + softmax_vars
  
      # all_vars = [self.emb, self.softmax_w, self.W, self.B, self.W_P, self.softmax_b]
      loss = self.forward(x, y, w, training=True)
      scaled_loss = loss * FLAGS.num_steps
      grads = tf.gradients(scaled_loss, all_vars)
  
      # Embedding grads
      emb_grads = grads[: len(emb_vars)]
      emb_grads = [
          tf.IndexedSlices(
              grad.values * FLAGS.batch_size,
              grad.indices,
              grad.dense_shape,
          )
          for grad in emb_grads
      ]
  
      lstm_grads = grads[len(emb_vars) : len(emb_vars) + len(lstm_vars)]
      lstm_grads, _ = tf.clip_by_global_norm(
          lstm_grads, FLAGS.max_grad_norm
      )
  
      softmax_grads = grads[len(emb_vars) + len(lstm_vars) :]
  
      clipped_grads = emb_grads + lstm_grads + softmax_grads
      grads_and_vars = list(zip(clipped_grads, all_vars))
  
      train_op = self.optimizer.apply_gradients(grads_and_vars)
      # FIXME(Hao): The code below does not work in TF2.x,
      # ema = tf.train.ExponentialMovingAverage(decay=0.999)
      # with tf.control_dependencies([train_op]):
      # ema.apply([self.softmax_w, self.softmax_b])
      return loss, train_op
  
file_path: autodist/examples/lm1b/language_model.py
imports: import os
  import numpy
  import tensorflow
  from absl import flags
label_desc: Use image node_id 3 for calling the LM obj's underlying member method code with example usage: obj.train_step(input_data) and returns: loss,  train_op