def __init__(self,
               n_heads=2, d_qk=64, d_v=64, share_qk='unused',
               causal=False,
               masked=False,
               chunk_len=128, n_chunks_before=1, n_chunks_after=0,
               n_hashes=1,
               n_buckets=None,
               mode='train',
               predict_mem_len=2048, predict_drop_len=256,
               attention_dropout=0.0,
               output_dropout=0.0,
               max_length_for_buckets=None,
               bias=False,
               n_parallel_heads=1,
               use_python_loop=False,
               use_reference_code=False,
              ):
    """Construct an LSH self-attention layer."""
    super().__init__(n_in=(2 if masked else 1), n_out=1)

    self._n_heads = n_heads
    if n_parallel_heads:
      if ((n_parallel_heads > n_heads and n_parallel_heads % n_heads != 0)
          or (n_parallel_heads < n_heads and n_heads % n_parallel_heads != 0)):
        raise ValueError(
            'n_parallel_heads must be a multiple or fraction of n_heads')
      self._n_parallel_heads = n_parallel_heads
    else:
      self._n_parallel_heads = None

    self._incremental = (mode == 'predict')
    if self._incremental:
      assert causal, 'Only causal attention supports fast inference'
      assert chunk_len is not None or (predict_mem_len and predict_drop_len)
      predict_mem_len = predict_mem_len or (chunk_len * (1 + n_chunks_before))
      predict_drop_len = predict_drop_len or chunk_len
      if predict_mem_len is None or predict_drop_len is None:
        raise ValueError('This configuration does not support fast inference.')
      if not 0 < predict_drop_len <= predict_mem_len:
        raise ValueError(
            'Bad parameter values: (predict_mem_len, predict_drop_len) = ',
            predict_mem_len, predict_drop_len)
      self._predict_mem_len = predict_mem_len
      self._predict_drop_len = predict_drop_len

    self._use_python_loop = use_python_loop
    self._use_reference_code = use_reference_code

    self._d_qk = d_qk
    self._d_v = d_v
    self._share_qk = True
    self._causal = causal
    self._masked = masked
    self._chunk_len = chunk_len
    self._n_chunks_before = n_chunks_before
    self._n_chunks_after = n_chunks_after
    self._bias = bias
    self._mode = mode
    if mode == 'train':
      self._attention_dropout = attention_dropout
      self._output_dropout = output_dropout
    else:
      self._attention_dropout = 0.0
      self._output_dropout = 0.0

    self._n_hashes = n_hashes
    self._n_buckets = n_buckets
    self._max_length_for_buckets = max_length_for_buckets
Construct an LSH self-attention layer.
trax/layers/research/efficient_attention.py