(
    q, k=None, v=None,
    q_chunk_len=None, kv_chunk_len=None,
    n_chunks_before=0, n_chunks_after=0,
    mask_fn=None, q_info=None, kv_info=None,
    dropout=0.0, rng=None,
    )
return out, dots_logsumexp
Dot-product attention, with optional chunking and/or masking.

  Args:
    q: Query vectors, shape [q_len, d_qk]
    k: Key vectors, shape [kv_len, d_qk]; or None
    v: Value vectors, shape [kv_len, d_v]
    q_chunk_len: Set to non-zero to enable chunking for query vectors
    kv_chunk_len: Set to non-zero to enable chunking for key/value vectors
    n_chunks_before: Number of adjacent previous chunks to attend to
    n_chunks_after: Number of adjacent subsequent chunks to attend to
    mask_fn: TODO(kitaev) doc
    q_info: Query-associated metadata for masking
    kv_info: Key-associated metadata for masking
    dropout: Dropout rate
    rng: RNG for dropout

  Returns:
    A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and
    dots_logsumexp has shape [q_len]. The logsumexp of the attention
    probabilities is useful for combining multiple rounds of attention (as in
    LSH attention).
Dot-product attention, with optional chunking and/or masking.
Dot - product attention with optional chunking and / or masking .
def attend(
    q, k=None, v=None,
    q_chunk_len=None, kv_chunk_len=None,
    n_chunks_before=0, n_chunks_after=0,
    mask_fn=None, q_info=None, kv_info=None,
    dropout=0.0, rng=None,
    ):
  """Dot-product attention, with optional chunking and/or masking.

  Args:
    q: Query vectors, shape [q_len, d_qk]
    k: Key vectors, shape [kv_len, d_qk]; or None
    v: Value vectors, shape [kv_len, d_v]
    q_chunk_len: Set to non-zero to enable chunking for query vectors
    kv_chunk_len: Set to non-zero to enable chunking for key/value vectors
    n_chunks_before: Number of adjacent previous chunks to attend to
    n_chunks_after: Number of adjacent subsequent chunks to attend to
    mask_fn: TODO(kitaev) doc
    q_info: Query-associated metadata for masking
    kv_info: Key-associated metadata for masking
    dropout: Dropout rate
    rng: RNG for dropout

  Returns:
    A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and
    dots_logsumexp has shape [q_len]. The logsumexp of the attention
    probabilities is useful for combining multiple rounds of attention (as in
    LSH attention).
  """
  assert v is not None
  share_qk = (k is None)

  # `q_info` and `kv_info` if supplied are 0 indexed, we want them to be 1
  # indexed instead so that we can mask position 0 as well - see Github #820

  if q_info is None:
    q_info = np.arange(1, q.shape[-2] + 1, dtype=np.int32)
  else:
    q_info += 1

  if kv_info is None and not share_qk:
    kv_info = np.arange(1, v.shape[-2] + 1, dtype=np.int32)
  elif kv_info is not None:
    kv_info += 1

  # Split q/k/v into chunks along the time axis, if desired.
  if q_chunk_len is not None:
    q = np.reshape(q, (-1, q_chunk_len, q.shape[-1]))
    q_info = np.reshape(q_info, (-1, q_chunk_len))

  if share_qk:
    assert kv_chunk_len is None or kv_chunk_len == q_chunk_len
    k = q
    kv_chunk_len = q_chunk_len
    if kv_info is None:
      kv_info = q_info
    elif kv_chunk_len is not None:
      # kv_info is not None, but reshape as required.
      kv_info = np.reshape(kv_info, (-1, kv_chunk_len))
  elif kv_chunk_len is not None:
    k = np.reshape(k, (-1, kv_chunk_len, k.shape[-1]))
    kv_info = np.reshape(kv_info, (-1, kv_chunk_len))

  if kv_chunk_len is not None:
    v = np.reshape(v, (-1, kv_chunk_len, v.shape[-1]))

  if share_qk:
    k = length_normalized(k)
  k = k / np.sqrt(k.shape[-1])

  # Optionally include adjacent chunks.
  if q_chunk_len is not None or kv_chunk_len is not None:
    assert q_chunk_len is not None and kv_chunk_len is not None
  else:
    assert n_chunks_before == 0 and n_chunks_after == 0

  k = look_adjacent(k, n_chunks_before, n_chunks_after)
  v = look_adjacent(v, n_chunks_before, n_chunks_after)
  kv_info = look_adjacent(kv_info, n_chunks_before, n_chunks_after)

  # Dot-product attention.
  dots = np.matmul(q, np.swapaxes(k, -1, -2))

  # Masking
  if mask_fn is not None:
    dots = mask_fn(dots, q_info[..., :, None], kv_info[..., None, :])

  # Softmax.
  dots_logsumexp = fastmath.logsumexp(dots, axis=-1, keepdims=True)
  dots = np.exp(dots - dots_logsumexp)

  if dropout > 0.0:
    assert rng is not None
    # Dropout is broadcast across the bin dimension
    dropout_shape = (dots.shape[-2], dots.shape[-1])
    # TODO(kitaev): verify that tie-in is safe to remove (in light of jax fix)
    keep_prob = 1.0 - dropout
    keep = fastmath.random.bernoulli(rng, keep_prob, dropout_shape)
    multiplier = keep.astype(dots.dtype) / keep_prob
    dots = dots * multiplier

  # The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn.
  out = np.matmul(dots, v)
  out = np.reshape(out, (-1, out.shape[-1]))
  dots_logsumexp = np.reshape(dots_logsumexp, (-1,))
  return out, dots_logsumexp
def attend ( q , k = None , v = None , q_chunk_len = None , kv_chunk_len = None , n_chunks_before = 0 , n_chunks_after = 0 , mask_fn = None , q_info = None , kv_info = None , dropout = 0.0 , rng = None , ) : assert v is not None share_qk = ( k is None ) # `q_info` and `kv_info` if supplied are 0 indexed, we want them to be 1 # indexed instead so that we can mask position 0 as well - see Github #820 if q_info is None : q_info = np . arange ( 1 , q . shape [ - 2 ] + 1 , dtype = np . int32 ) else : q_info += 1 if kv_info is None and not share_qk : kv_info = np . arange ( 1 , v . shape [ - 2 ] + 1 , dtype = np . int32 ) elif kv_info is not None : kv_info += 1 # Split q/k/v into chunks along the time axis, if desired. if q_chunk_len is not None : q = np . reshape ( q , ( - 1 , q_chunk_len , q . shape [ - 1 ] ) ) q_info = np . reshape ( q_info , ( - 1 , q_chunk_len ) ) if share_qk : assert kv_chunk_len is None or kv_chunk_len == q_chunk_len k = q kv_chunk_len = q_chunk_len if kv_info is None : kv_info = q_info elif kv_chunk_len is not None : # kv_info is not None, but reshape as required. kv_info = np . reshape ( kv_info , ( - 1 , kv_chunk_len ) ) elif kv_chunk_len is not None : k = np . reshape ( k , ( - 1 , kv_chunk_len , k . shape [ - 1 ] ) ) kv_info = np . reshape ( kv_info , ( - 1 , kv_chunk_len ) ) if kv_chunk_len is not None : v = np . reshape ( v , ( - 1 , kv_chunk_len , v . shape [ - 1 ] ) ) if share_qk : k = length_normalized ( k ) k = k / np . sqrt ( k . shape [ - 1 ] ) # Optionally include adjacent chunks. if q_chunk_len is not None or kv_chunk_len is not None : assert q_chunk_len is not None and kv_chunk_len is not None else : assert n_chunks_before == 0 and n_chunks_after == 0 k = look_adjacent ( k , n_chunks_before , n_chunks_after ) v = look_adjacent ( v , n_chunks_before , n_chunks_after ) kv_info = look_adjacent ( kv_info , n_chunks_before , n_chunks_after ) # Dot-product attention. dots = np . matmul ( q , np . swapaxes ( k , - 1 , - 2 ) ) # Masking if mask_fn is not None : dots = mask_fn ( dots , q_info [ ... , : , None ] , kv_info [ ... , None , : ] ) # Softmax. dots_logsumexp = fastmath . logsumexp ( dots , axis = - 1 , keepdims = True ) dots = np . exp ( dots - dots_logsumexp ) if dropout > 0.0 : assert rng is not None # Dropout is broadcast across the bin dimension dropout_shape = ( dots . shape [ - 2 ] , dots . shape [ - 1 ] ) # TODO(kitaev): verify that tie-in is safe to remove (in light of jax fix) keep_prob = 1.0 - dropout keep = fastmath . random . bernoulli ( rng , keep_prob , dropout_shape ) multiplier = keep . astype ( dots . dtype ) / keep_prob dots = dots * multiplier # The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn. out = np . matmul ( dots , v ) out = np . reshape ( out , ( - 1 , out . shape [ - 1 ] ) ) dots_logsumexp = np . reshape ( dots_logsumexp , ( - 1 , ) ) return out , dots_logsumexp