code: def EncoderDecoderBlock(
      d_model,
      d_ff,
      n_heads,
      dropout,
      dropout_shared_axes,
      mode,
      ff_activation,
      ff_dropout,
      ff_chunk_size,
      ff_use_sru,
      ff_sparsity,
      ff_sparsity_type,
      attention_chunk_size,
      attention_type,
      enc_dec_attention_sparsity=0,
  ):
      """Returns a list of layers implementing a Transformer encoder-decoder block.
  
      The input is a triple (decoder_activations, mask, encoder_activiations) where
      the mask is created from the original input token IDs to prevent attending to
      the padding part of the encoder.
  
      Args:
        d_model: Final dimension of tensors at most points in the model, including
          the initial embedding output.
        d_ff: Size of special dense layer in the feed-forward part of each block.
        n_heads: Number of attention heads.
        dropout: Stochastic rate (probability) for dropping an activation value when
          applying dropout within a block.
        dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
          along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
          way to save memory and apply consistent masks to activation vectors at
          different sequence positions.
        mode: If `'train'`, each block will include dropout; else, it will pass all
          values through unaltered.
        ff_activation: Type of activation function at the end of each block; must be
          an activation-type subclass of `Layer`.
        ff_dropout: Stochastic rate (probability) for dropping an activation value
          when applying dropout after the FF dense layer.
        ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
        ff_use_sru: int or pair of ints; if > 0, we use this many SRU layers
          in addition to the feed-forward block (second int specifies sru size)
        ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
         ff_sparsity_type: string, if ff_sparsity >0,
          use SparseFF if ff_sparsity_type=`'1inN'` and
          use BlockSparseFF if ff_sparsity_type=`'Block'`
        attention_chunk_size: int, if > 0 run attention chunked at this size
        attention_type: The attention layer to use.
        enc_dec_attention_sparsity: Sparsity to use in encoder-decoder attention.
  
      Returns:
        A list of layers which maps triples (decoder_activations, mask,
        encoder_activations) to triples of the same sort.
      """
  
      def _Dropout():
          return tl.Dropout(
              rate=dropout, shared_axes=dropout_shared_axes, mode=mode
          )
  
      # TODO(afrozm): This layer isn't configurable because: We currently don't have
      # any alternative for it (LSH cannot do it fundamentally, that's why we have
      # NoEncDec models, and local attention doesn't make sense in the general
      # setting where we don't know what in input is local to what in output;
      # some variants of FAVOR can do it, so maybe in the future,
      # but we don't have them yet).
      if isinstance(enc_dec_attention_sparsity, tuple):
          q_sparsity, result_sparsity = enc_dec_attention_sparsity
      elif enc_dec_attention_sparsity > 0:
          q_sparsity = enc_dec_attention_sparsity
          result_sparsity = (
              "noop"  # We simply skip Dense layer after attention.
          )
      else:
          q_sparsity = None
          result_sparsity = None
      attention_qkv = tl.AttentionQKV(
          d_model,
          n_heads=n_heads,
          dropout=dropout,
          mode=mode,
          cache_KV_in_predict=True,
          q_sparsity=q_sparsity,
          result_sparsity=result_sparsity,
      )
  
      causal_attention = ApplyAttentionLayer(
          attention_type,
          d_model,
          n_heads,
          d_model // n_heads,
          d_model // n_heads,
          causal=True,
          masked=True,
          attention_dropout=dropout,
          output_dropout=dropout,
          attention_chunk_size=attention_chunk_size,
          mode=mode,
      )
  
      feed_forward = FeedForwardWithOptions(
          d_model,
          d_ff,
          dropout,
          dropout_shared_axes,
          ff_activation,
          ff_dropout,
          ff_chunk_size,
          ff_use_sru,
          ff_sparsity,
          True,
          mode,
          False,
          ff_sparsity_type,
      )
  
      return [  # vec_d masks vec_e
          tl.Residual(
              tl.LayerNorm(),  # vec_d ..... .....
              causal_attention,  # vec_d ..... .....
              _Dropout(),  # vec_d ..... .....
          ),
          tl.Residual(
              tl.LayerNorm(),  # vec_d ..... .....
              tl.Select(
                  [0, 2, 2, 1, 2]
              ),  # vec_d vec_e vec_e masks vec_e
              attention_qkv,  # vec_d masks vec_e
              _Dropout(),  # vec_d masks vec_e
          ),
          tl.Residual(feed_forward),  # vec_d masks vec_e
      ]
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 11 for calling a global function with example usage: EncoderDecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes, mode, ff_activation, ff_dropout, ff_chunk_size, ff_use_sru, ff_sparsity, ff_sparsity_type, attention_chunk_size, attention_type, enc_dec_attention_sparsity) and returns: list, tl