code: def ChunkedDecoderLayer(feature_depth,
                          feedforward_depth,
                          num_heads,
                          dropout,
                          chunk_selector,
                          mode):
    
    return layers.Serial(
        layers.Residual(  # Self-attention block.
            layers.Map(layers.LayerNorm()),
            layers.ChunkedCausalMultiHeadedAttention(
                feature_depth, num_heads=num_heads, dropout=dropout,
                chunk_selector=chunk_selector, mode=mode),
            layers.Map(layers.Dropout(rate=dropout, mode=mode)),
        ),
        layers.Map(ResidualFeedForward(
            feature_depth, feedforward_depth, dropout, mode=mode))
    )
docs: Transformer decoder layer operating on chunks.
  
    Args:
      feature_depth: int:  depth of embedding
      feedforward_depth: int: depth of feed-forward layer
      num_heads: int: number of attention heads
      dropout: float: dropout rate (how much to drop out)
      chunk_selector: a function from chunk number to list of chunks to attend.
      mode: str: 'train' or 'eval'
  
    Returns:
      the layer.