code: def get_sequence_lengths(sequence_batch: Array, eos_id: int) -> Array:
      """Returns the length of each one-hot sequence, including the EOS token."""
      # sequence_batch.shape = (batch_size, seq_length, vocab_size)
      eos_row = sequence_batch[:, :, eos_id]
      eos_idx = jnp.argmax(eos_row, axis=-1)  # returns first occurrence
      # `eos_idx` is 0 if EOS is not present, so we use full length in that case.
      return jnp.where(
          eos_row[jnp.arange(eos_row.shape[0]), eos_idx],
          eos_idx + 1,
          sequence_batch.shape[
              1
          ],  # if there is no EOS, use full length
      )
  
imports: import random
  from typing import Any, Dict, Generator, Optional, Tuple
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: get_sequence_lengths(sequence_batch, eos_id) and returns: jnp