def EncoderDecoderMask():
  """Returns a layer that creates a mask for encoder-decoder cross attention.

  The layer expects two inputs:

      - decoder_input: batch of integer (e.g., token ID) sequences
      - mask: padding mask from the encoder

  The layer output is a mask that marks for each sequence position (for both
  encoder and decoder) whether that position can be attended to or not. The
  encoder-decoder mask shape is (batch_size, 1, decoder_sequence_length,
  encoder_sequence_length), such that axis 1 will automatically broadcast to
  cover any number of attention heads.
  """
  def f(decoder_input, mask):
    if len(decoder_input.shape) != 3:
      raise ValueError(
          f'Decoder input to EncoderDecoderMask must be a 3-D array with '
          f'shape (batch_size, decoder_sequence_length, d_model); instead got '
          f'shape {decoder_input.shape}.')
    batch_size = mask.shape[0]
    encoder_sequence_length = mask.shape[-1]
    decoder_sequence_length = decoder_input.shape[1]
    mask = mask.reshape((batch_size, 1, 1, encoder_sequence_length))
    return mask + jnp.zeros((1, 1, decoder_sequence_length, 1))
  return Fn('EncoderDecoderMask', f)
Returns a layer that creates a mask for encoder-decoder cross attention.

  The layer expects two inputs:

      - decoder_input: batch of integer (e.g., token ID) sequences
      - mask: padding mask from the encoder

  The layer output is a mask that marks for each sequence position (for both
  encoder and decoder) whether that position can be attended to or not. The
  encoder-decoder mask shape is (batch_size, 1, decoder_sequence_length,
  encoder_sequence_length), such that axis 1 will automatically broadcast to
  cover any number of attention heads.
  
trax/layers/attention.py