def __call__(self, inputs_q, inputs_kv, mask=None):
    """Apply the Transformer layer.

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

    Returns:
      layer_output: <float32>[batch_size, len, hidden_size].
    """
    attn_output = TransformerAttention(
        hidden_size=self.hidden_size,
        attention_dropout_rate=self.attention_dropout_rate,
        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,
        num_heads=self.num_heads,
        name='attention')(inputs_q, inputs_kv, mask)
    intermediate_output = TransformerIntermediate(
        intermediate_size=self.intermediate_size,
        hidden_act=self.hidden_act,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        name='intermediate')(
            attn_output)
    layer_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')(intermediate_output, attn_output)
    return layer_output
Apply the Transformer layer.

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

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