def __call__(self, x: named_axes.NamedArray) -> named_axes.NamedArray:
    """Runs the attention computation.

    Args:
      x: The input to the computation, which will be mapped to queries, keys,
        and values by the sublayers.

    Returns:
      The final output of the ``attn_value_to_output`` sublayer.
    """
    query = self.input_to_query(x)
    key = self.input_to_key(x)
    value = self.input_to_value(x)
    attn = self.query_key_to_attn((query, key))
    output = self.attn_value_to_output((attn, value))
    return output
Runs the attention computation.

    Args:
      x: The input to the computation, which will be mapped to queries, keys,
        and values by the sublayers.

    Returns:
      The final output of the ``attn_value_to_output`` sublayer.
    
penzai/deprecated/v1/nn/attention.py