code: 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])
  
imports: from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
  from trax.models.research import configurable_transformer
label_desc: Use image node_id 3 for calling a global function with example usage: _StripFromConcatenateWithPadding(vec_ed, tok_e, tok_d) and returns: fastmath, fastmath