code: def ApplyAttentionLayer(
      attention_type,
      d_model,
      n_heads,
      d_qk,
      d_v,
      causal,
      masked,
      attention_dropout,
      output_dropout,
      attention_chunk_size,
      mode,
  ):
      """Runs the supplied attention layer."""
      try:
          attention = attention_type(
              n_heads=n_heads,
              d_qk=d_qk,
              d_v=d_v,
              causal=causal,
              masked=masked,
              output_dropout=output_dropout,
              attention_dropout=attention_dropout,
              mode=mode,
          )
      except TypeError:  # No d_qk arguments in less advanced layers.
          attention = attention_type(
              d_model,
              n_heads=n_heads,
              dropout=attention_dropout,
              mode=mode,
          )
      return tl.Chunk(attention, attention_chunk_size)
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 3 for calling a global function with example usage: ApplyAttentionLayer(attention_type, d_model, n_heads, d_qk, d_v, causal, masked, attention_dropout, output_dropout, attention_chunk_size, mode) and returns: tl