def __call__(self, inputs_q, inputs_kv, mask=None):
    """Apply attention block including attention and an output layer.

    Args:
      inputs_q: <float32>[batch_size, len, dim].
      inputs_kv: <float32>[batch_size, len, dim].
      mask: <int32>[batch_size, len].

    Returns:
      output: <float32>[batch_size, len, hidden_size].
    """
    output = nn.MultiHeadDotProductAttention(
        qkv_features=self.hidden_size,
        num_heads=self.num_heads,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        dropout_rate=self.attention_dropout_rate,
        decode=self.decode,
        name='attention',
    )(inputs_q, inputs_kv, mask=mask, deterministic=self.deterministic)
    output = TransformerOutput(
        hidden_size=self.hidden_size,
        hidden_dropout_rate=self.hidden_dropout_rate,
        layer_norm_eps=self.layer_norm_eps,
        deterministic=self.deterministic,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='output')(output, inputs_q)
    return output
Apply attention block including attention and an output layer.

    Args:
      inputs_q: <float32>[batch_size, len, dim].
      inputs_kv: <float32>[batch_size, len, dim].
      mask: <int32>[batch_size, len].

    Returns:
      output: <float32>[batch_size, len, hidden_size].
    
language/gscan/xattn_model/model/layers.py