class_docstr: Sequence-to-sequence class using encoder/decoder architecture.
  
  Attributes:
    teacher_force: whether to use `decoder_inputs` as input to the decoder at
      every step. If False, only the first input (i.e., the "=" token) is used,
      followed by samples taken from the previous output logits.
    hidden_size: int, the number of hidden dimensions in the encoder and decoder
      LSTMs.
    vocab_size: the size of the vocabulary.
    eos_id: EOS id.
class_docstr_tok: ["Sequence-to-sequence","class","using","encoder\/decoder","architecture",".","Attributes",":","teacher_force",":","whether","to","use","`","decoder_inputs","`","as","input","to","the","decoder","at","every","step",".","If","False",",","only","the","first","input","(","i.e.",",","the","``","=","''","token",")","is","used",",","followed","by","samples","taken","from","the","previous","output","logits",".","hidden_size",":","int",",","the","number","of","hidden","dimensions","in","the","encoder","and","decoder","LSTMs",".","vocab_size",":","the","size","of","the","vocabulary",".","eos_id",":","EOS","id","."]
docstr: Applies the seq2seq model.
  
  Args:
    encoder_inputs: [batch_size, max_input_length, vocab_size].
      padded batch of input sequences to encode.
    decoder_inputs: [batch_size, max_output_length, vocab_size].
      padded batch of expected decoded sequences for teacher forcing.
      When sampling (i.e., `teacher_force = False`), only the first token is
      input into the decoder (which is the token "="), and samples are used
      for the following inputs. The second dimension of this tensor determines
      how many steps will be decoded, regardless of the value of
      `teacher_force`.
  
  Returns:
    Pair (logits, predictions), which are two arrays of length `batch_size`
    containing respectively decoded logits and predictions (in one hot
    encoding format).
docstr_tok: ["Applies","the","seq2seq","model",".","Args",":","encoder_inputs",":","[","batch_size",",","max_input_length",",","vocab_size","]",".","padded","batch","of","input","sequences","to","encode",".","decoder_inputs",":","[","batch_size",",","max_output_length",",","vocab_size","]",".","padded","batch","of","expected","decoded","sequences","for","teacher","forcing",".","When","sampling","(","i.e.",",","`","teacher_force","=","False","`",")",",","only","the","first","token","is","input","into","the","decoder","(","which","is","the","token","``","=","''",")",",","and","samples","are","used","for","the","following","inputs",".","The","second","dimension","of","this","tensor","determines","how","many","steps","will","be","decoded",",","regardless","of","the","value","of","`","teacher_force","`",".","Returns",":","Pair","(","logits",",","predictions",")",",","which","are","two","arrays","of","length","`","batch_size","`","containing","respectively","decoded","logits","and","predictions","(","in","one","hot","encoding","format",")","."]
code: def __call__(
      self, encoder_inputs: Array, decoder_inputs: Array
  ) -> Tuple[Array, Array]:
      """Applies the seq2seq model.
  
      Args:
        encoder_inputs: [batch_size, max_input_length, vocab_size].
          padded batch of input sequences to encode.
        decoder_inputs: [batch_size, max_output_length, vocab_size].
          padded batch of expected decoded sequences for teacher forcing.
          When sampling (i.e., `teacher_force = False`), only the first token is
          input into the decoder (which is the token "="), and samples are used
          for the following inputs. The second dimension of this tensor determines
          how many steps will be decoded, regardless of the value of
          `teacher_force`.
  
      Returns:
        Pair (logits, predictions), which are two arrays of length `batch_size`
        containing respectively decoded logits and predictions (in one hot
        encoding format).
      """
      # Encode inputs.
      encoder = nn.RNN(
          nn.LSTMCell(self.hidden_size),
          return_carry=True,
          name="encoder",
      )
      decoder = nn.RNN(
          DecoderLSTMCell(
              decoder_inputs.shape[-1],
              self.teacher_force,
              self.vocab_size,
          ),
          split_rngs={"params": False, "lstm": True},
          name="decoder",
      )
  
      seq_lengths = self.get_seq_lengths(encoder_inputs)
  
      encoder_state, _ = encoder(
          encoder_inputs, seq_lengths=seq_lengths
      )
      logits, predictions = decoder(
          decoder_inputs[:, :-1],
          initial_carry=(encoder_state, decoder_inputs[:, 0]),
      )
  
      return logits, predictions
  
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 Seq2seq obj's underlying member method code with example usage: obj.__call__(encoder_inputs, decoder_inputs) and returns: logits,  predictions