def _StripFromConcatenateWithPadding(vec_ed, tok_e, tok_d):
  """Strip concatenate with padding: see the layer below for details."""
  # pylint: disable=invalid-name
  B, L, H = vec_ed.shape
  L1 = tok_e.shape[1]
  L2 = tok_d.shape[1]
  # pylint: enable=invalid-name
  if L != L1 + L2:
    raise ValueError(f'Length from encoder-decoder vectors ({L}) does not'
                     f' equal sum of lengths from encoder ({L1}) and decoder'
                     f' ({L2}).')
  if tok_e.shape != (B, L1):
    raise ValueError(f'Shape of encoder tokens, {tok_e.shape}, does not'
                     f' equal {(B, L1)}.')
  if tok_d.shape != (B, L2):
    raise ValueError(f'Shape of decoder tokens, {tok_d.shape}, does not'
                     f' equal {(B, L2)}.')

  def _UpdateRow(x):
    # (L, H), (L1, H) & (L2, H)
    row_ed, row_e, _ = x
    mask_e = row_e != 0
    len_e = jnp.sum(mask_e, dtype=jnp.int32)
    # In `row_ed` start where encoder tokens/vecs end, i.e. are index `len_e`
    # and pick up (L2, H) tensor slice from there.
    zero = jnp.array(0, dtype=len_e.dtype)  # avoid int32/int64 mismatch
    return fastmath.dynamic_slice(row_ed, (len_e, zero), (L2, H))

  return fastmath.map(_UpdateRow, [vec_ed, tok_e, tok_d])
Strip concatenate with padding: see the layer below for details.
trax/models/research/transformer2.py