def __init__(self,
               n_heads=2, d_qk=64, d_v=64, share_qk=False,
               causal=False, masked=False,
               chunk_len=None, n_chunks_before=0, n_chunks_after=0,
               bias=False,
               mode='train',
               predict_mem_len=None, predict_drop_len=None,
               attention_dropout=0.0,
               output_dropout=0.0,
               n_parallel_heads=None,
               use_python_loop=False,
               use_reference_code=False,
              ):
    """Construct a self-attention layer.

    Args:
      n_heads: int: Number of attention heads
      d_qk: int: Depth of query ond key vectors
      d_v: int: Depth of value vectors
      share_qk: bool: Set to True to share query and key projection weights
      causal: bool: Set to True to mask out attention to future items
      masked: bool: Set to True to accept an additional mask argument, that
        allows masking out attention to padding tokens.
      chunk_len (optional): Number of tokens per chunk. Setting this option will
        enable chunked attention.
      n_chunks_before: Number of previous chunks to attend to, when using
        chunked attention.
      n_chunks_after: Number of subsequent chunks to attend to, when using
        chunked attention. Don't use this option for causal attention, because
        attention to future tokens will be masked out anyway. However, note that
        cross-chunk attention "wraps around" in both directions, so this option
        is never a strict no-op.
      bias: bool: Set to True to add bias vectors when computing query/key/value
      mode: 'train', 'eval', or 'predict'
      predict_mem_len: int: Number of input positions to remember in a cache
        when doing fast inference. Whenever the cache fills up, some input
        elements will be forgotten. When chunking is enabled, the default is to
        store chunk_len * (1 + n_chunks_before) elements.
      predict_drop_len: int: Number of input elements to drop once the fast
        inference input cache fills up. When chunking is enabled, the default is
        to drop exactly chunk_len elements.
      attention_dropout: Dropout probability for attention mask.
      output_dropout: Dropout probability for the layer output.
      n_parallel_heads: Number of attention heads to compute in parallel.

          - If `n_parallel_heads` is None (default), the entire layer is
            computed with maximum parallelism. This mode is the fastest, but
            also uses the most memory. Start with this mode, but switch to one
            of the others if memory runs out.
          - If `n_parallel_heads` is 1, attention is computed one head at a
            time, and one example at a time. This mode uses the least memory
            but is not as fast as batched attention. Use this mode when working
            with very long sequences, such that any amount of parallelism won't
            fit in memory.
          - If `n_parallel_heads` is a multiple of `n_heads`, attention is
            computed for sub-batches of (`n_parallel_heads // n_heads`)
            examples at a time.
          - If `1 < n_parallel_heads < n_heads`, attention is computed for
            several heads at a time, but only within a single example. It must
            be the case that `n_heads` is a multiple of `n_parallel_heads`. Use
            this mode for long sequences, to strike a balance between
            parallelism and memory usage.
      use_python_loop: Set to True to use a Python loop when iterating over
          sub-batches of examples/heads (as opposed to a JAX/XLA loop).
          This option will increase compilation time and jitted code size,
          potentially drastically. Using it is not recommended except for
          testing/debugging. In particular, note that enabling this option on
          TPU can decrease the maximum model size that will fit in memory.
      use_reference_code: Set to True to fall back to the reference
          implementation of batched attention. This option will increase
          compilation time and jitted code size, potentially drastically. Using
          it is not recommended except for testing/debugging.

    """
    super().__init__(n_in=(2 if masked else 1), n_out=1)

    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._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._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 = share_qk
    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
Construct a self-attention layer.

    Args:
      n_heads: int: Number of attention heads
      d_qk: int: Depth of query ond key vectors
      d_v: int: Depth of value vectors
      share_qk: bool: Set to True to share query and key projection weights
      causal: bool: Set to True to mask out attention to future items
      masked: bool: Set to True to accept an additional mask argument, that
        allows masking out attention to padding tokens.
      chunk_len (optional): Number of tokens per chunk. Setting this option will
        enable chunked attention.
      n_chunks_before: Number of previous chunks to attend to, when using
        chunked attention.
      n_chunks_after: Number of subsequent chunks to attend to, when using
        chunked attention. Don't use this option for causal attention, because
        attention to future tokens will be masked out anyway. However, note that
        cross-chunk attention "wraps around" in both directions, so this option
        is never a strict no-op.
      bias: bool: Set to True to add bias vectors when computing query/key/value
      mode: 'train', 'eval', or 'predict'
      predict_mem_len: int: Number of input positions to remember in a cache
        when doing fast inference. Whenever the cache fills up, some input
        elements will be forgotten. When chunking is enabled, the default is to
        store chunk_len * (1 + n_chunks_before) elements.
      predict_drop_len: int: Number of input elements to drop once the fast
        inference input cache fills up. When chunking is enabled, the default is
        to drop exactly chunk_len elements.
      attention_dropout: Dropout probability for attention mask.
      output_dropout: Dropout probability for the layer output.
      n_parallel_heads: Number of attention heads to compute in parallel.

          - If `n_parallel_heads` is None (default), the entire layer is
            computed with maximum parallelism. This mode is the fastest, but
            also uses the most memory. Start with this mode, but switch to one
            of the others if memory runs out.
          - If `n_parallel_heads` is 1, attention is computed one head at a
            time, and one example at a time. This mode uses the least memory
            but is not as fast as batched attention. Use this mode when working
            with very long sequences, such that any amount of parallelism won't
            fit in memory.
          - If `n_parallel_heads` is a multiple of `n_heads`, attention is
            computed for sub-batches of (`n_parallel_heads // n_heads`)
            examples at a time.
          - If `1 < n_parallel_heads < n_heads`, attention is computed for
            several heads at a time, but only within a single example. It must
            be the case that `n_heads` is a multiple of `n_parallel_heads`. Use
            this mode for long sequences, to strike a balance between
            parallelism and memory usage.
      use_python_loop: Set to True to use a Python loop when iterating over
          sub-batches of examples/heads (as opposed to a JAX/XLA loop).
          This option will increase compilation time and jitted code size,
          potentially drastically. Using it is not recommended except for
          testing/debugging. In particular, note that enabling this option on
          TPU can decrease the maximum model size that will fit in memory.
      use_reference_code: Set to True to fall back to the reference
          implementation of batched attention. This option will increase
          compilation time and jitted code size, potentially drastically. Using
          it is not recommended except for testing/debugging.

    
trax/layers/research/efficient_attention.py