code: def sequence_mask(lengths: Array, max_length: int) -> Array:
      """Computes a boolean mask over sequence positions for each given length.
  
      Example:
      ```
      sequence_mask([1, 2], 3)
      [[True, False, False],
       [True, True, False]]
      ```
  
      Args:
        lengths: The length of each sequence. <int>[batch_size]
        max_length: The width of the boolean mask. Must be >= max(lengths).
  
      Returns:
        A mask with shape: <bool>[batch_size, max_length] indicating which
        positions are valid for each sequence.
      """
      return jnp.arange(max_length)[None] < lengths[:, None]
  
imports: import functools
  from typing import Any, Callable, Optional
  from flax import linen
  import jax
  from jax import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: sequence_mask(lengths, max_length) and returns: unknown