def _fast_inference_update_state(inputs, state, mask_for_predict=None):
  """Updates state of a causal attention layer for fast inference.

  The layer state stores arrays with cached values of keys and values,
  as well as an index. To make shapes static, keys and values in the state are
  long, and the index indicates where the new keys and values from inputs need
  to be appended.

  During update, we append new_keys and new_values to keys and values at
  position given by index. And we increment index by length of new keys.
  We also create a mask to be 1 at appropriate positions (causal mask).

  Args:
    inputs: a triple (new_queries, new_keys, new_values)
    state: layer state with (keys, values, index)
    mask_for_predict: mask used for predict mode. This is used only in
      Terraformer.

  Returns:
    Updated state and mask to be used.
  """
  # Fast inference: run step-by-step, storing the sequence
  # of keys and values calculated so far in state.
  (_, new_k, new_v) = inputs
  if mask_for_predict is not None:
    (state_mask_for_predict, ks, vs, idx) = state
  else:
    (ks, vs, idx) = state
  length = new_k.shape[1]
  # TODO(lukaszkaiser): benchmark speed and decide if using a separate code path
  # with index_update when length == 1 is worth it.
  # Keys and values are of shape [batch_size, length, d_kv].
  ks = fastmath.dynamic_update_slice_in_dim(ks, new_k, idx, axis=1)
  vs = fastmath.dynamic_update_slice_in_dim(vs, new_v, idx, axis=1)
  k_length = ks.shape[1]

  # Mask is of shape [1, q_length, k_length].
  # Mask should be true for every pair of (query_token, key_token) such that
  # index of query_token is equal or larger to index of key_token.
  mask = (jnp.reshape(jnp.arange(k_length), (1, 1, k_length))
          <= jnp.reshape(jnp.arange(length) + idx, (1, length, 1)))
  if mask_for_predict is None:
    return (ks, vs, idx + length), mask
  else:
    state_mask_for_predict = fastmath.dynamic_update_slice_in_dim(
        state_mask_for_predict != 0, mask_for_predict.reshape((-1)) != 0, 0,
        axis=0)

    state_mask_for_predict = fastmath.dynamic_update_slice_in_dim(
        state_mask_for_predict != 0, jnp.ones((1,)) != 0,
        jnp.sum(mask_for_predict, dtype=jnp.int32), axis=0)

    state_mask_for_predict = fastmath.dynamic_update_slice_in_dim(
        state_mask_for_predict != 0, jnp.ones((1,)) != 0, idx, axis=0)
    placeholder = jnp.reshape(state_mask_for_predict != 0,
                              (1, 1, mask.shape[2],))
    mask = mask * placeholder

    return (state_mask_for_predict, ks, vs, idx + length), mask
Updates state of a causal attention layer for fast inference.

  The layer state stores arrays with cached values of keys and values,
  as well as an index. To make shapes static, keys and values in the state are
  long, and the index indicates where the new keys and values from inputs need
  to be appended.

  During update, we append new_keys and new_values to keys and values at
  position given by index. And we increment index by length of new keys.
  We also create a mask to be 1 at appropriate positions (causal mask).

  Args:
    inputs: a triple (new_queries, new_keys, new_values)
    state: layer state with (keys, values, index)
    mask_for_predict: mask used for predict mode. This is used only in
      Terraformer.

  Returns:
    Updated state and mask to be used.
  
trax/layers/attention.py