def MultiplicativeConvCausalAttention(
    d_feature, n_heads=1, sparsity=None, length_kernel_size=3, dropout=0.0,
    force_no_dropout=False, max_inference_length=2048, share_qk=False,
    output_layer_type='none', v_concat_type='none', mode='train'):
  """Returns a layer that maps activations to activations, with causal masking.

  Like `CausalAttention`, this layer type represents one pass of multi-head
  self-attention with causal masking rather than padding-based masking. However,
  for computing Q/K/V instead of a Dense layer it combines
  FactoredDense layer with LocallyConvLayer.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    sparsity: The sparsity of the layer; usually it should be equal to n_heads.
    length_kernel_size: Size of convolution kernel on the length dimension.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    force_no_dropout: If True, force dropout to be 0.0 independent of the above
        value; used to override some configurations.
    max_inference_length: maximum length for inference.
    share_qk: if True, average Q and K embeddings and share for both Q and K.
    output_layer_type: Which sparse layers to use for processing output from the
        attention mechanism. One of `'none'`, `'mult'`, `'conv'`,
        or `'multconv'`.
    v_concat_type: What kind of concatenation to use when computing V tensor.
        One of `'original'`, `'fixed'`, or `'none'`. `'none'` means using just
        output from mutliplicative layer shared by Q, K, V. `'fixed'` means
        using output from multiplicative layer concatenated, for each module,
        with the layer input. `'original'` means using concatenation without
        properly taking modules into account; this method was used in
        experiments previously, so it is included for backwards-compatibility.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
  """
  assert output_layer_type in ['none', 'mult', 'conv', 'multconv']
  assert v_concat_type in ['original', 'fixed', 'none']

  dropout = 0.0 if force_no_dropout else dropout
  sparsity = n_heads if sparsity is None else sparsity
  d_module = d_feature // sparsity

  output_layers = []
  if 'mult' in output_layer_type:
    output_layers.append(FactoredDense(
        sparsity, d_feature, d_feature))
  if 'conv' in output_layer_type:
    output_layers.append(LocallyConvDense(
        sparsity, d_module, mode=mode, kernel_size=3,
        length_kernel_size=length_kernel_size))

  if v_concat_type == 'original':
    # 'original'` uses concatenation without properly taking modules into
    # account; this method was used in experiments previously, so it is included
    # for backwards-compatibility.
    concat_layers = [tl.Concatenate()]  # use permuted and original for v
  elif v_concat_type == 'fixed':
    # `'fixed'` uses the output from multiplicative layer concatenated, for each
    # module, with the layer input. This means that every module in Conv layer
    # has access both to parts of embeddings which were used to compute Q/K of
    # this particular module, and it ha access to parts of the embedding which
    # will be modified by this module.
    concat_layers = [
        tl.Parallel(
            tl.Fn('Reshape1', lambda x: jnp.reshape(  # pylint: disable=g-long-lambda
                x, (x.shape[0], x.shape[1], sparsity, d_module))),
            tl.Fn('Reshape2', lambda x: jnp.reshape(  # pylint: disable=g-long-lambda
                x, (x.shape[0], x.shape[1], sparsity, d_module)))),
        tl.Concatenate(),
        tl.Fn('Reshape3',
              lambda x: jnp.reshape(x, (x.shape[0], x.shape[1], 2*d_feature))),
    ]
  elif v_concat_type == 'none':
    # `'none'` doesn't use concatenation: we throw away the original layer
    # input and pass to Conv only output of shared Multiplicative layer.
    concat_layers = [tl.Select([0], n_in=2)]

  if share_qk:
    return tl.Serial(
        tl.Select([0, 0]),  # pre-qkv, pre-v-for-concat
        FactoredDense(sparsity, d_feature, d_feature),  # shared q k
        tl.Select([0, 0]),  # pre-qk, pre-v, pre-v-for-concat
        LocallyConvDense(sparsity, d_module, mode=mode, kernel_size=3,
                         length_kernel_size=length_kernel_size),
        tl.SplitIntoHeads(n_heads),
        tl.Select([0, 0]),  # use for q and k
        tl.Parallel(
            [],
            [],
            [concat_layers,
             LocallyConvDense(sparsity, d_module, mode=mode, kernel_size=1,
                              length_kernel_size=length_kernel_size),
             tl.SplitIntoHeads(n_heads)],
        ),
        tl.DotProductCausalAttention(
            dropout=dropout, max_inference_length=max_inference_length,
            mode=mode),
        tl.MergeHeads(n_heads),
        output_layers,
    )
  return tl.Serial(
      tl.Select([0, 0]),  # duplicate activations
      FactoredDense(sparsity, d_feature, d_feature),  # shared q, k
      tl.Select([0, 0, 0]),  # use for q, k, v
      tl.Parallel(
          [LocallyConvDense(sparsity, d_module, mode=mode, kernel_size=3,
                            length_kernel_size=length_kernel_size),
           tl.SplitIntoHeads(n_heads)],
          [LocallyConvDense(sparsity, d_module, mode=mode, kernel_size=3,
                            length_kernel_size=length_kernel_size),
           tl.SplitIntoHeads(n_heads)],
          [concat_layers,
           LocallyConvDense(sparsity, d_module, mode=mode, kernel_size=1,
                            length_kernel_size=length_kernel_size),
           tl.SplitIntoHeads(n_heads)],
      ),
      tl.DotProductCausalAttention(
          dropout=dropout, max_inference_length=max_inference_length,
          mode=mode),
      tl.MergeHeads(n_heads),
      output_layers,
  )
Returns a layer that maps activations to activations, with causal masking.

  Like `CausalAttention`, this layer type represents one pass of multi-head
  self-attention with causal masking rather than padding-based masking. However,
  for computing Q/K/V instead of a Dense layer it combines
  FactoredDense layer with LocallyConvLayer.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    sparsity: The sparsity of the layer; usually it should be equal to n_heads.
    length_kernel_size: Size of convolution kernel on the length dimension.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    force_no_dropout: If True, force dropout to be 0.0 independent of the above
        value; used to override some configurations.
    max_inference_length: maximum length for inference.
    share_qk: if True, average Q and K embeddings and share for both Q and K.
    output_layer_type: Which sparse layers to use for processing output from the
        attention mechanism. One of `'none'`, `'mult'`, `'conv'`,
        or `'multconv'`.
    v_concat_type: What kind of concatenation to use when computing V tensor.
        One of `'original'`, `'fixed'`, or `'none'`. `'none'` means using just
        output from mutliplicative layer shared by Q, K, V. `'fixed'` means
        using output from multiplicative layer concatenated, for each module,
        with the layer input. `'original'` means using concatenation without
        properly taking modules into account; this method was used in
        experiments previously, so it is included for backwards-compatibility.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
  
trax/layers/research/sparsity.py