Class TransformerLayerGenerate

transformer_layer.py:55–527

Full transformer layer, with attention.

@gin.configurable
class TransformerLayerGenerate(transformer_layer.TransformerLayer):
  """Full transformer layer, with attention."""

  def _next_decoder_state(
      self, decoder_state: DecoderState, keys: Array, values: Array
  ) -> Tuple[DecoderState, Array, Array]:
    """Compute the next decoder state, and return keys,values to attend to.

    The keys,values returned from this function are drawn from the prior
    decoding state, and comprise a full window of local context.

    Args:
      decoder_state: The current decoder state, initially created using
        init_decoder_state().
      keys: The key for the current token, of shape (batch_size, 1, dim)
      values: The value for the current token of shape (batch_size, 1, dim)

    Returns:
      (next_decoder_state,
       window of keys of shape (batch_size, window_length, dim),
       window of values of shape (batch_size, window_length, dim))
    """

    assert keys.shape[1] == 1  # single-token autoregressive decoding.

    # Unpack decoder_state
    stored_keys = decoder_state["keys"]
    stored_values = decoder_state["values"]
    curr_index = decoder_state["current_index"]

    # Slice to get window_length-sized chunk of previous keys,values.
    out_decoder_state = {}
    curr_win_index = curr_index - self.window_length

    # out_keys = jax.lax.dynamic_slice_in_dim(
    #     stored_keys, curr_win_index, self.window_length, axis=1)
    out_keys = slice_in_dim_1(self.window_length)(stored_keys, curr_win_index)

    # out_values = jax.lax.dynamic_slice_in_dim(
    #     stored_values, curr_win_index, self.window_length, axis=1)
    out_values = slice_in_dim_1(self.window_length)(
        stored_values, curr_win_index
    )

    # Write current keys,values to stored keys, values.
    # stored_keys = jax.lax.dynamic_update_slice_in_dim(
    #     stored_keys, keys, curr_index, axis=1)
    stored_keys = update_slice_in_dim_1(stored_keys, keys, curr_index)
    # stored_values = jax.lax.dynamic_update_slice_in_dim(
    #     stored_values, values, curr_index, axis=1)
    stored_values = update_slice_in_dim_1(stored_values, values, curr_index)
    curr_index = curr_index + 1

    # Pack a new decoder_state object.
    out_decoder_state["keys"] = stored_keys
    out_decoder_state["values"] = stored_values
    out_decoder_state["current_index"] = curr_index
    out_decoder_state["relative_position_bias"] = decoder_state[