code: 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),  # 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))
    )
docs: 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.