class_docstr: DecoderLSTM Module wrapped in a lifted scan transform.
  
  Attributes:
    teacher_force: See docstring on Seq2seq module.
    vocab_size: Size of the vocabulary.
class_docstr_tok: ["DecoderLSTM","Module","wrapped","in","a","lifted","scan","transform",".","Attributes",":","teacher_force",":","See","docstring","on","Seq2seq","module",".","vocab_size",":","Size","of","the","vocabulary","."]
docstr: Applies the DecoderLSTM model.
docstr_tok: ["Applies","the","DecoderLSTM","model","."]
code: def __call__(
      self, carry: Tuple[LSTMCarry, Array], x: Array
  ) -> Tuple[Tuple[LSTMCarry, Array], Tuple[Array, Array]]:
      """Applies the DecoderLSTM model."""
      lstm_state, last_prediction = carry
      if not self.teacher_force:
          x = last_prediction
      lstm_state, y = nn.LSTMCell(self.features)(lstm_state, x)
      logits = nn.Dense(features=self.vocab_size)(y)
      # Sample the predicted token using a categorical distribution over the
      # logits.
      categorical_rng = self.make_rng("lstm")
      predicted_token = jax.random.categorical(categorical_rng, logits)
      # Convert to one-hot encoding.
      prediction = jax.nn.one_hot(
          predicted_token, self.vocab_size, dtype=jnp.float32
      )
  
      return (lstm_state, prediction), (logits, prediction)
  
file_path: flax/examples/seq2seq/models.py
imports: from typing import Tuple
  from flax import linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the DecoderLSTMCell obj's underlying member method code with example usage: obj.__call__(carry, x) and returns: lstm_state,  prediction,  logits,  prediction