def __call__(self, input1, input2, mask1=None, mask2=None):
    """Apply the Transformer connection 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:
      The output representation for input1 and input2 with shape
      (batch_size, len, hidden_size).
    """
    attn_output1, attn_output2 = TransformerCrossAttention(
        num_heads=self.bi_num_heads,
        bi_hidden_size=self.bi_hidden_size,
        hidden_size1=self.hidden_size1,
        hidden_size2=self.hidden_size2,
        attention_dropout_rate=self.attention_dropout_rate,
        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_attention')(input1, input2, mask1, mask2)
    intermediate_layer = functools.partial(
        TransformerIntermediate,
        hidden_act=self.hidden_act,
        kernel_init=self.kernel_init,
        bias_init=self.bias_init)
    output_layer = functools.partial(
        TransformerOutput,
        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)
    intermediate_output1 = intermediate_layer(
        intermediate_size=self.intermediate_size1, name='intermediate1')(
            attn_output1)
    layer_output1 = output_layer(
        hidden_size=self.hidden_size1, name='output1')(intermediate_output1,
                                                       attn_output1)
    intermediate_output2 = intermediate_layer(
        intermediate_size=self.intermediate_size2, name='intermediate2')(
            attn_output2)
    layer_output2 = output_layer(
        hidden_size=self.hidden_size2, name='output2')(intermediate_output2,
                                                       attn_output2)
    return layer_output1, layer_output2
Apply the Transformer connection 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:
      The output representation for input1 and input2 with shape
      (batch_size, len, hidden_size).
    
language/gscan/xattn_model/model/layers.py