code: def MultiHeadedAttention(
      feature_depth, num_heads=8, dropout=0.0, mode='train'):
    
    return combinators.Serial(
        combinators.Parallel(
            combinators.Branch(num_branches=3),  # q = k = v = first input
            combinators.Identity()  # pass the mask
        ),
        MultiHeadedAttentionQKV(  # pylint: disable=no-value-for-parameter
            feature_depth, num_heads=num_heads, dropout=dropout, mode=mode),
    )
docs: Transformer-style multi-headed attention.
  
    Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
  
    Args:
      feature_depth: int:  depth of embedding
      num_heads: int: number of attention heads
      dropout: float: dropout rate
      mode: str: 'train' or 'eval'
  
    Returns:
      Multi-headed self-attention layer.