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: Get segmentation mask for inputs.
docstr_tok: ["Get","segmentation","mask","for","inputs","."]
code: def get_seq_lengths(self, inputs: Array) -> Array:
      """Get segmentation mask for inputs."""
      # undo one-hot encoding
      inputs = jnp.argmax(inputs, axis=-1)
      # calculate sequence lengths
      seq_lengths = jnp.argmax(inputs == self.eos_id, axis=-1)
  
      return seq_lengths
  
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 2 for calling the Seq2seq obj's underlying member method code with example usage: obj.get_seq_lengths(inputs) and returns: seq_lengths