def __call__(self, x, encoded, decoder_mask=None, encoder_decoder_mask=None):
    """Apply the transfomer encoder-decoder layer.

    Args:
      x: <float32>[batch_size, len,  dim].
      encoded: <float32>[batch_size, len, dim].
      decoder_mask: <int32>[batch_size, len].
      encoder_decoder_mask: <int32>[batch_size, len].

    Returns:
      layer_output: <float32>[batch_size, len, hidden_size].
    """
    transformer_attention = functools.partial(
        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)
    attn_output = transformer_attention(
        name='attention', decode=self.decode)(
            x, x, mask=decoder_mask)
    attn_output = transformer_attention(name='encoder_decoder_attention')(
        attn_output, encoded, mask=encoder_decoder_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 transfomer encoder-decoder layer.

    Args:
      x: <float32>[batch_size, len,  dim].
      encoded: <float32>[batch_size, len, dim].
      decoder_mask: <int32>[batch_size, len].
      encoder_decoder_mask: <int32>[batch_size, len].

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