class_docstr: TokenLearner module using MHA.
  
  Attributes:
    num_tokens: Number of tokens to generate.
    num_heads: Number of heads to use for the dot product attention.
docstr: Applies TokenLearner-mha to the inputs.
  
  Args:
    inputs: Inputs of shape `[bs, hw, c]`.
  
  Returns:
    Output of shape `[bs, num_tokens, c]`.
docstr_tok: ["Applies","TokenLearner-mha","to","the","inputs",".","Args",":","inputs",":","Inputs","of","shape","`","[","bs",",","hw",",","c","]","`",".","Returns",":","Output","of","shape","`","[","bs",",","num_tokens",",","c","]","`","."]
code: def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
      """Applies TokenLearner-mha to the inputs.
  
      Args:
        inputs: Inputs of shape `[bs, hw, c]`.
  
      Returns:
        Output of shape `[bs, num_tokens, c]`.
      """
      bs, _, d = inputs.shape
  
      token_init = nn.initializers.normal(stddev=0.02)
      tokens = self.param("tokens", token_init, (1, self.num_tokens, d))
      tokens = jnp.broadcast_to(tokens, (bs, self.num_tokens, d))
      return nn.MultiHeadDotProductAttention(
          num_heads=self.num_heads, deterministic=True, name="attn"
      )(inputs_q=tokens, inputs_kv=inputs)
  
imports: from typing import Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import mixer
  from scenic.projects.baselines import vit
  from scenic.projects.token_learner import model
label_desc: Use image node_id 1 for calling the TokenLearnerMHA obj's underlying member method code with example usage: obj.__call__(inputs) and returns: nn