code: def mask_sequences(sequence_batch: Array, lengths: Array) -> Array:
      """Sets positions beyond the length of each sequence to 0."""
      return sequence_batch * (
          lengths[:, np.newaxis]
          > np.arange(sequence_batch.shape[1])[np.newaxis]
      )
  
imports: import random
  from typing import Any, Dict, Generator, Optional, Tuple
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: mask_sequences(sequence_batch, lengths) and returns: unknown