def _ConcatWithPadding(vec_e, vec_d, mask_e):
  """Concatenate with padding: see the ConcatWithPadding layer for details."""
  # pylint: disable=invalid-name
  B, L1, H = vec_e.shape
  L2 = vec_d.shape[1]
  # pylint: enable=invalid-name

  if vec_d.shape != (B, L2, H):
    raise ValueError(f'Shape of decoder vector, {vec_d.shape}, does not'
                     f' equal {(B, L2, H)}.')
  if mask_e.shape != (B, L1):
    raise ValueError(f'Shape of encoder mask, {mask_e.shape}, does not'
                     f' equal {(B, L1)}.')

  def _UpdateRow(x):
    # row_e - (L1, H), row_d - (L2, H), row_mask_e - (L1,)
    row_e, row_d, row_mask_e = x
    # final_row - (L1+L2, H)
    final_row = jnp.concatenate([row_e, jnp.zeros_like(row_d)], axis=0)
    # Find the last real token/vector of the encoder.
    e_idx = jnp.sum(row_mask_e, dtype=jnp.int32)
    # Starting after that index, update with the decoder row.
    zero = jnp.array(0, dtype=e_idx.dtype)  # avoid int32/int64 mismatch
    return fastmath.dynamic_update_slice(final_row, row_d, (e_idx, zero))

  return fastmath.map(_UpdateRow, [vec_e, vec_d, mask_e])
Concatenate with padding: see the ConcatWithPadding layer for details.
trax/models/research/transformer2.py