def sequence_mask(lengths, maxlen=None, dtype=tf.bool, transpose=False):
  """Returns a sequence mask, like tf.sequence_mask().

  Unlike tf.sequence_mask(), this can also generate a transposed mask, which is
  convenient for working with time-major sequences.

  Args:
    lengths: <int>[batch_size] Sequence lengths.
    maxlen: <int>[] Maximum length, or None to compute the max of `lengths`.
    dtype: DType of the generated mask.
    transpose: Whether to generate the transpose of the mask.

  Returns:
    <dtype>[maxlen, batch_size] Sequence mask.
  """
  with tf.name_scope("sequence_mask"):
    if maxlen is None:
      maxlen = tf.reduce_max(lengths)
    positions = tf.range(maxlen, dtype=lengths.dtype)

    positions_singleton_axis = 1 if transpose else 0
    lengths_singleton_axis = 0 if transpose else 1
    positions = tf.expand_dims(positions, axis=positions_singleton_axis)
    lengths = tf.expand_dims(lengths, axis=lengths_singleton_axis)
    mask = positions < lengths

    if dtype != tf.bool:
      mask = tf.cast(mask, dtype)
    return mask
Returns a sequence mask, like tf.sequence_mask().

  Unlike tf.sequence_mask(), this can also generate a transposed mask, which is
  convenient for working with time-major sequences.

  Args:
    lengths: <int>[batch_size] Sequence lengths.
    maxlen: <int>[] Maximum length, or None to compute the max of `lengths`.
    dtype: DType of the generated mask.
    transpose: Whether to generate the transpose of the mask.

  Returns:
    <dtype>[maxlen, batch_size] Sequence mask.
  
language/common/utils/tensor_utils.py