class_docstr: Computes MLP-based attention scores based on keys alone, without a query.
  
  Attention scores are computed by feeding the keys through an MLP. This
  results in a single scalar per key, and for each sequence the attention
  scores are normalized using a softmax so that they sum to 1. Invalid key
  positions are ignored as indicated by the mask. This is also called
  "Bahdanau attention" and was originally proposed in:
  ```
  Bahdanau et al., 2015. Neural Machine Translation by Jointly Learning to
  Align and Translate. ICLR. https://arxiv.org/abs/1409.0473
  ```
  
  Attributes:
    hidden_size: The hidden size of the MLP that computes the attention score.
docstr: Applies model to the input keys and mask.
  
  Args:
    keys: The inputs for which to compute an attention score. Shape:
      <float32>[batch_size, seq_length, embeddings_size].
    mask: A mask that determines which values in `keys` are valid. Only
      values for which the mask is True will get non-zero attention scores.
      <bool>[batch_size, seq_length].
  
  Returns:
    The normalized attention scores. <float32>[batch_size, seq_length].
code: def __call__(self, keys: Array, mask: Array) -> Array:
      """Applies model to the input keys and mask.
  
      Args:
        keys: The inputs for which to compute an attention score. Shape:
          <float32>[batch_size, seq_length, embeddings_size].
        mask: A mask that determines which values in `keys` are valid. Only
          values for which the mask is True will get non-zero attention scores.
          <bool>[batch_size, seq_length].
  
      Returns:
        The normalized attention scores. <float32>[batch_size, seq_length].
      """
      hidden = nn.Dense(self.hidden_size, name="keys", use_bias=False)(
          keys
      )
      energy = nn.tanh(hidden)
      scores = nn.Dense(1, name="energy", use_bias=False)(energy)
      scores = scores.squeeze(
          -1
      )  # New shape: <float32>[batch_size, seq_len].
      scores = jnp.where(
          mask, scores, -jnp.inf
      )  # Using exp(-inf) = 0 below.
      scores = nn.softmax(scores, axis=-1)
  
      # Captures the scores if 'intermediates' is mutable, otherwise does nothing.
      self.sow("intermediates", "attention", scores)
  
      return scores
  
imports: import functools
  from typing import Any, Callable, Optional
  from flax import linen
  import jax
  from jax import numpy
label_desc: Use image node_id 1 for calling the KeysOnlyMlpAttention obj's underlying member method code with example usage: obj.__call__(keys, mask) and returns: scores