class_docstr: Attention pooling layer.
  
  Attributes:
    num_heads: Number of heads.
    features: Number of features.
class_docstr_tok: ["Attention","pooling","layer",".","Attributes",":","num_heads",":","Number","of","heads",".","features",":","Number","of","features","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      x = x.reshape(x.shape[0], -1, x.shape[3])
  
      x = jnp.concatenate([x.mean(axis=1, keepdims=True), x], axis=1)
  
      positional_embedding = self.param(
          "positional_embedding",
          jax.nn.initializers.normal(1.0 / x.shape[-1] ** 0.5),
          (x.shape[1], x.shape[2]),
      )
      attn = nn.MultiHeadDotProductAttention(
          self.num_heads,
          qkv_features=x.shape[-1],
          use_bias=True,
          out_features=self.features,
          name="attn",
      )
  
      x = x + positional_embedding[jnp.newaxis].astype(x.dtype)
      x = attn(x[:, :1], x)
      return x[:, 0]
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the AttentionPool obj's underlying member method code with example usage: obj.__call__(x) and returns: x