code: def test_mask_sequences(self):
      np.testing.assert_equal(
          input_pipeline.mask_sequences(
              np.arange(1, 13).reshape((4, 3)), np.array([3, 2, 1, 0])
          ),
          np.array([[1, 2, 3], [4, 5, 0], [7, 0, 0], [0, 0, 0]]),
      )
  
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 2 for calling the TrainTest obj's underlying member method code with example usage: obj.test_mask_sequences() without return types