def __call__(self, input1, input2, mask1=None, mask2=None):
    """Apply attention block including cross-attention and an output layer.

    Args:
      input1: <float32>[batch_size, len, dim].
      input2: <float32>[batch_size, len, dim].
      mask1: <int32>[batch_size, len].
      mask2: <int32>[batch_size, len].

    Returns:
      output: <float32>[batch_size, len, hidden_size].
    """
    context1 = nn.MultiHeadDotProductAttention(
        qkv_features=self.bi_hidden_size,
        out_features=self.bi_hidden_size,
        num_heads=self.num_heads,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        dropout_rate=self.attention_dropout_rate,
        name='cross1')(
            input2, input1, mask=mask1, deterministic=self.deterministic)
    context2 = nn.MultiHeadDotProductAttention(
        qkv_features=self.bi_hidden_size,
        out_features=self.bi_hidden_size,
        num_heads=self.num_heads,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        dropout_rate=self.attention_dropout_rate,
        name='cross2')(
            input1, input2, mask=mask2, deterministic=self.deterministic)
    output1, output2 = TransformerBiOutput(
        hidden_size1=self.hidden_size1,
        hidden_size2=self.hidden_size2,
        hidden_dropout_rate=self.hidden_dropout_rate,
        layer_norm_eps=self.layer_norm_eps,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init,
        deterministic=self.deterministic,
        name='bi_output')(context2, input1, context1, input2)
    return output1, output2
Apply attention block including cross-attention and an output layer.

    Args:
      input1: <float32>[batch_size, len, dim].
      input2: <float32>[batch_size, len, dim].
      mask1: <int32>[batch_size, len].
      mask2: <int32>[batch_size, len].

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