def padding_mask(seq_k, seq_q):
  """For masking out the padding part of the keys sequence.

  Args:
    seq_k: Keys tensor, with shape [B, L_k]
    seq_q: Query tensor, with shape [B, L_q]

  Returns:
    A masking tensor, with shape [B, L_1, L_k]
  """
  len_q = seq_q.size(1)
  # `PAD` is 0
  pad_mask = seq_k.eq(0)
  pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1)  # shape [B, L_q, L_k]
  return pad_mask