(vec_e, vec_d, mask_e)
return fastmath.map(_UpdateRow, [vec_e, vec_d, mask_e])
Concatenate with padding: see the ConcatWithPadding layer for details.
Concatenate with padding: see the ConcatWithPadding layer for details.
Concatenate with padding : see the ConcatWithPadding layer for details .
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])
def _ConcatWithPadding ( vec_e , vec_d , mask_e ) : # 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 ] )