def ChunkedCausalMultiHeadedAttention(
    feature_depth, num_heads=8, dropout=0.0, chunk_selector=None, mode='train'):
  """Transformer-style causal multi-headed attention operating on chunks.

  Accepts inputs that are a list of chunks and applies causal attention.

  Args:
    feature_depth: int:  depth of embedding
    num_heads: int: number of attention heads
    dropout: float: dropout rate
    chunk_selector: a function from chunk number to list of chunks to attend.
    mode: str: 'train' or 'eval'

  Returns:
    Multi-headed self-attention layer.
  """
  prepare_attention_input = combinators.Serial(
      combinators.Branch(),
      combinators.Parallel(
          combinators.Branch(num_branches=3),  # q = k = v = first input
          CausalMask(axis=-2),  # pylint: disable=no-value-for-parameter
      ),
      combinators.Parallel(
          combinators.Parallel(
              core.Dense(feature_depth),
              core.Dense(feature_depth),
              core.Dense(feature_depth),
          ),
          combinators.Identity()
      )
  )
  return combinators.Serial(
      combinators.Map(prepare_attention_input),
      ChunkedAttentionSelector(selector=chunk_selector),  # pylint: disable=no-value-for-parameter
      combinators.Map(PureMultiHeadedAttention(  # pylint: disable=no-value-for-parameter
          feature_depth=feature_depth, num_heads=num_heads,
          dropout=dropout, mode=mode), check_shapes=False),
      combinators.Map(core.Dense(feature_depth))
  )
def ChunkedCausalMultiHeadedAttention ( feature_depth , num_heads = 8 , dropout = 0.0 , chunk_selector = None , mode = 'train' ) : prepare_attention_input = combinators . Serial ( combinators . Branch ( ) , combinators . Parallel ( combinators . Branch ( num_branches = 3 ) , CausalMask ( axis = - 2 ) , ) , combinators . Parallel ( combinators . Parallel ( core . Dense ( feature_depth ) , core . Dense ( feature_depth ) , core . Dense ( feature_depth ) , ) , combinators . Identity ( ) ) ) return combinators . Serial ( combinators . Map ( prepare_attention_input ) , ChunkedAttentionSelector ( selector = chunk_selector ) , combinators . Map ( PureMultiHeadedAttention ( feature_depth = feature_depth , num_heads = num_heads , dropout = dropout , mode = mode ) , check_shapes = False ) , combinators . Map ( core . Dense ( feature_depth ) ) )
Transformer-style causal multi-headed attention operating on chunks.

  Accepts inputs that are a list of chunks and applies causal attention.

  Args:
    feature_depth: int:  depth of embedding
    num_heads: int: number of attention heads
    dropout: float: dropout rate
    chunk_selector: a function from chunk number to list of chunks to attend.
    mode: str: 'train' or 'eval'

  Returns:
    Multi-headed self-attention layer.