def look_adjacent(x, n_chunks_before, n_chunks_after):
  """Used to implement attention between consecutive chunks.

  Args:
    x: array of shape [n_chunks, chunk_len, ...]
    n_chunks_before: Number of previous chunks to attend to.
    n_chunks_after: Number of subsequent chunks to attend to.
  Returns:
    array of shape [n_chunks, N * chunk_len, ...], where
    N = (1 + n_chunks_before + n_chunks_after).
  """
  if n_chunks_before == 0 and n_chunks_after == 0:
    return x

  slices = []
  for i in range(-n_chunks_before, n_chunks_after + 1):
    if i == 0:
      slices.append(x)
    else:
      slices.append(np.concatenate([x[i:, ...], x[:i, ...]], axis=0))
  return np.concatenate(slices, axis=1)
Used to implement attention between consecutive chunks.

  Args:
    x: array of shape [n_chunks, chunk_len, ...]
    n_chunks_before: Number of previous chunks to attend to.
    n_chunks_after: Number of subsequent chunks to attend to.
  Returns:
    array of shape [n_chunks, N * chunk_len, ...], where
    N = (1 + n_chunks_before + n_chunks_after).
  
trax/layers/research/efficient_attention.py