code: def forward(self, x, y, w, training):
      # [bs, steps, emb_size]
      x = tf.compat.v1.nn.embedding_lookup(
          self.emb, x, partition_strategy="div"
      )
      if training:
          x = tf.nn.dropout(x, 1 - self.keep_prob)
  
      # [bs, emb_size] * steps
      inputs = [
          tf.squeeze(v, axis=[1])
          for v in tf.split(
              value=x, num_or_size_splits=self.num_steps, axis=1
          )
      ]
      for t in range(self.num_steps):
          cell_inputs = tf.concat([inputs[t], self.h], axis=1)
          lstm_matrix = tf.linalg.matmul(cell_inputs, self.W) + self.B
          i, j, f, o = tf.split(lstm_matrix, 4, axis=1)
          self.c = tf.math.sigmoid(f + 1.0) * self.c + tf.math.sigmoid(
              i
          ) * tf.math.tanh(j)
          self.h = tf.math.sigmoid(o) * tf.math.tanh(self.c)
          self.h = tf.linalg.matmul(self.h, self.W_P)
          inputs[t] = self.h
          if training:
              inputs[t] = tf.nn.dropout(inputs[t], 1 - self.keep_prob)
  
      inputs[t] = tf.identity(inputs[t])
      inputs = tf.reshape(
          tf.concat(inputs, axis=1), [-1, self.projected_size]
      )
  
      if training:
          targets = tf.reshape(y, [-1, 1])
          loss = tf.nn.sampled_softmax_loss(
              self.softmax_w,
              self.softmax_b,
              targets,
              inputs,
              self.num_sampled,
              self.vocab_size,
          )
      else:
          full_softmax_w = tf.reshape(
              tf.concat(self.softmax_w, axis=1),
              [-1, self.projected_size],
          )
          full_softmax_w = full_softmax_w[: self.vocab_size, :]
  
          logits = (
              tf.matmul(inputs, full_softmax_w, transpose_b=True)
              + self.softmax_b
          )
          targets = tf.reshape(y, [-1])
          loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
              logits=logits, labels=targets
          )
  
      loss = tf.reduce_mean(
          loss * tf.reshape(tf.dtypes.cast(w, tf.float32), [-1])
      )
      return loss
  
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 2 for calling the LM obj's underlying member method code with example usage: obj.forward(x, y, w, training) and returns: loss