Create an attention mask to hide padding and future words.

def MakeTargetMask(target, pad=0):
  """Create an attention mask to hide padding and future words."""
  target_mask = (target != pad)[ :, np.newaxis, :]
  target_dtype = target_mask.dtype
  causal_mask = onp.tril(onp.ones((1, target.shape[-1], target.shape[-1]),
                                  dtype=target_dtype), k=0)
  target_mask = target_mask & causal_mask
  return np.expand_dims(target_mask, axis=1)