def PaddingMask(pad=0):
  """Returns a layer that maps integer sequences to padding masks.

  The layer expects as input a batch of integer sequences. The layer output is
  an N-D array that marks for each sequence position whether the integer (e.g.,
  a token ID) in that position represents padding -- value ``pad`` -- versus
  text/content -- all other values. The padding mask shape is
  (batch_size, 1, 1, encoder_sequence_length), such that axis 1 will broadcast
  to cover any number of attention heads and axis 2 will broadcast to cover
  decoder sequence positions.

  Args:
    pad: Integer that represents padding rather than a token/content ID.
  """
  def f(x):
    if len(x.shape) != 2:
      raise ValueError(
          f'Input to PaddingMask must be a 2-D array with shape '
          f'(batch_size, sequence_length); instead got shape {x.shape}.')
    batch_size = x.shape[0]
    sequence_length = x.shape[1]
    content_positions = (x != pad)
    return content_positions.reshape((batch_size, 1, 1, sequence_length))
  return Fn(f'PaddingMask({pad})', f)
Returns a layer that maps integer sequences to padding masks.

  The layer expects as input a batch of integer sequences. The layer output is
  an N-D array that marks for each sequence position whether the integer (e.g.,
  a token ID) in that position represents padding -- value ``pad`` -- versus
  text/content -- all other values. The padding mask shape is
  (batch_size, 1, 1, encoder_sequence_length), such that axis 1 will broadcast
  to cover any number of attention heads and axis 2 will broadcast to cover
  decoder sequence positions.

  Args:
    pad: Integer that represents padding rather than a token/content ID.
  
trax/layers/attention.py