code: def CreateAttentionMaskLayer():
      """Creates attention mask layer.
  
      Returns a layer that based on queries, keys and accumulated pool size of
      keys/values until this layer calculates positional embeddings for
      causal relative attention calculations.
  
      Takes as input q, k, v and appends proper mask in the end.
      Causal attention uses masking to prevent a given sequence position from
      attending to positions greater than / following it. This is used, for
      example, when training autoregressive sequence models, or when decoding a
      sequence symbol by symbol.
  
      Returns:
        an attention mask layer.
      """
  
      def calculate_mask(queries, keys):
          batch_size = queries.shape[0]
          keys_len, queries_len = keys.shape[-2], queries.shape[-2]
          funnel_factor, is_upsampling = calc_funnel_ratio(
              keys_len, queries_len
          )
  
          return _funnel_mask(
              batch_size,
              keys_len,
              queries_len,
              funnel_factor,
              is_upsampling,
          )
  
      def _funnel_mask(
          batch_size,
          keys_len,
          queries_len,
          funnel_factor,
          is_upsampling,
      ):
          """Funnel mask.
  
          Args:
            batch_size: batch size.
            keys_len: keys length.
            queries_len: queries length.
            funnel_factor: funnel factor.
            is_upsampling: True or False.
  
          Returns:
            funnel mask.
  
          This function based on keys/queries lengths creates a triangle mask
          that prevents tokens from attending to positions following it.
  
          If funnel_factor is not equal to 1 due to funnel upsampling or
          downsampling it adjusts created mask for funnel attention
          by repeating each element funnel_factor times.
  
          This is because after funnel layer one token attends to funnel_factor
          different tokens in downsampling. During upsampling on the other hand
          funnel_factor tokens are attending to single token before upsampling.
          """
  
          if funnel_factor != 1:
              if not is_upsampling:
                  mask = jnp.tril(
                      jnp.ones(
                          (queries_len, queries_len), dtype=jnp.bool_
                      )
                  )
                  mask = jnp.repeat(mask, funnel_factor, axis=-1)
              else:
                  mask = jnp.tril(
                      jnp.ones((keys_len, keys_len), dtype=jnp.bool_)
                  )
                  mask = jnp.repeat(mask, funnel_factor, axis=-2)
          else:
              mask = jnp.tril(
                  jnp.ones((queries_len, queries_len), dtype=jnp.bool_)
              )
  
          return jnp.repeat(mask[None, None, :, :], batch_size, axis=0)
  
      return cb.Branch(
          cb.Select([0]),
          cb.Select([1]),
          cb.Select([2]),
          cb.Fn("create attention mask layer", calculate_mask, n_out=1),
      )
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 9 for calling a global function with example usage: CreateAttentionMaskLayer() and returns: cb, _funnel_mask, jnp