code: 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])
  
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 2 for calling a global function with example usage: _ConcatWithPadding(vec_e, vec_d, mask_e) and returns: fastmath, fastmath