code: def test_get_sequence_lengths(self):
      oh_sequence_batch = jax.vmap(
          functools.partial(jax.nn.one_hot, num_classes=4)
      )(np.array([[0, 1, 0], [1, 0, 2], [1, 2, 0], [1, 2, 3]]))
      np.testing.assert_equal(
          input_pipeline.get_sequence_lengths(
              oh_sequence_batch, eos_id=0
          ),
          np.array([1, 2, 3, 3], np.int32),
      )
      np.testing.assert_equal(
          input_pipeline.get_sequence_lengths(
              oh_sequence_batch, eos_id=1
          ),
          np.array([2, 1, 1, 1], np.int32),
      )
      np.testing.assert_equal(
          input_pipeline.get_sequence_lengths(
              oh_sequence_batch, eos_id=2
          ),
          np.array([3, 3, 2, 2], np.int32),
      )
  
file_path: flax/examples/seq2seq/train_test.py
imports: import functools
  from absl.testing import absltest
  from flax.training import train_state
  import jax
  from jax import random
  import numpy
  import optax
  import input_pipeline
  import train
  import models
label_desc: Use image node_id 3 for calling the TrainTest obj's underlying member method code with example usage: obj.test_get_sequence_lengths() without return types