def __call__(self, hidden1, x1, hidden2, x2):
    """Apply the Transformer bi-output layer.

    Args:
      hidden1: <float32>[batch_size, len, dim].
      x1: <float32>[batch_size, len, dim].
      hidden2: <float32>[batch_size, len, dim].
      x2: <float32>[batch_size, len, dim].

    Returns:
      output: <float32>[batch_size, len, hidden_size1], <float32>[batch_size,
      len, hidden_size2]
    """
    transformer_output = functools.partial(
        TransformerOutput,
        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)
    hidden1 = transformer_output(
        hidden_size=self.hidden_size1, name='output1')(hidden1, x1)
    hidden2 = transformer_output(
        hidden_size=self.hidden_size2, name='output2')(hidden2, x2)
    return hidden1, hidden2
Apply the Transformer bi-output layer.

    Args:
      hidden1: <float32>[batch_size, len, dim].
      x1: <float32>[batch_size, len, dim].
      hidden2: <float32>[batch_size, len, dim].
      x2: <float32>[batch_size, len, dim].

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