class_docstr: Multihead Attention Pooling.
class_docstr_tok: ["Multihead","Attention","Pooling","."]
code: def __call__(self, x):
      n, l, d = x.shape  # pylint: disable=unused-variable
      probe = self.param(
          "probe", nn.initializers.xavier_uniform(), (1, 1, d), x.dtype
      )
      probe = jnp.tile(probe, [n, 1, 1])
  
      x = nn.MultiHeadDotProductAttention(
          num_heads=self.num_heads,
          kernel_init=nn.initializers.xavier_uniform(),
      )(probe, x)
  
      y = nn.LayerNorm()(x)
      x = x + MlpBlock(mlp_dim=self.mlp_dim)(y)
      return x[:, 0]
  
imports: import functools
  from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 1 for calling the MAPHead obj's underlying member method code with example usage: obj.__call__(x) and returns: x