Project: 013_Attention_as_a_Hypernetwork | Function: DotProductAttention.__call__ (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** RMSHead normalization implementation that normalizes attention weights across heads instead of keys, supporting the hypernetwork latent code interpretation

**KEY TERMS:** {}

**CONTEXT CODE:** """
Copyright (c) Simon Schug
All rights reserved.

MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
import functools
import warnings
from typing import Any, Callable, Optional, Tuple, Union, overload

import jax
import jax.numpy as jnp
from flax.linen import initializers
from flax.linen.attention import combine_masks
from flax.linen.dtypes import promote_dtype
from flax.linen.linear import DenseGeneral, PrecisionLike, default_kernel_init
from flax.linen.module import Module, compact, merge_param
from flax.linen.normalization import LayerNorm, RMSNorm
from jax import lax, random

PRNGKey = jax.Array
Shape = Tuple[int, ...]
Dtype = Any
Array = Any


class DotProductAttention(Module):
    normalization: str = "softmax"

    @compact
    def __call__(
        self,
        query: Array,
        key: Array,
        bias: Optional[Array] = None,
        mask: Optional[Array] = None,
        broadcast_dropout: bool = True,
        dropout_rng: Optional[PRNGKey] = None,
        dropout_rate: float = 0.0,
        deterministic: bool = False,
        dtype: Optional[Dtype] = None,
        precision: PrecisionLike = None,
    ):


        """This is the core function for applying attention based on
        https://arxiv.org/abs/1706.03762. It calculates the attention weights given
        query and key and combines the values using the attention weights.

        Note: query, key, value do not need to have any batch dimensions.

        Args:
        query: queries for calculating attention with shape of `[batch..., q_length,
            num_heads, qk_depth_per_head]`.
        key: keys for calculating attention with shape of `[batch..., kv_length,
            num_heads, qk_depth_per_head]`.
        value: values to be used in attention with shape of `[batch..., kv_length,
            num_heads, v_depth_per_head]`.
        bias: bias for the attention weights. This should be broadcastable to the
            shape `[batch..., num_heads, q_length, kv_length]`. This can be used for
            incorporating causal masks, padding masks, proximity bias, etc.
        mask: mask for the attention weights. This should be broadcastable to the
            shape `[batch..., num_heads, q_length, kv_length]`. This can be used for
            incorporating causal masks. Attention weights are masked out if their
            corresponding mask value is `False`.
        broadcast_dropout: bool: use a broadcasted dropout along batch dims.
        dropout_rng: JAX PRNGKey: to be used for dropout
        dropout_rate: dropout rate
        deterministic: bool, deterministic or not (to apply dropout)
        dtype: the dtype of the computation (default: infer from inputs)
        precision: numerical precision of the computation see `jax.lax.Precision`
            for details.

        Returns:
        Output of shape `[batch..., q_length, num_heads, v_depth_per_head]`.
        """
        raise NotImplementedError("Function implementation has been masked for testing")


class MultiHeadDotProductAttention(Module):
    """Multi-head dot-product attention block.

    Attributes:
      num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
        should be divisible by the number of heads.
      dtype: the dtype of the computation (default: infer from inputs and params)
      param_dtype: the dtype passed to parameter initializers (default: float32)
      qk_features: dimension of the key, query.
      v_features: dimension of the value.
      out_features: dimension of the last projection
      broadcast_dropout: bool: use a broadcasted dropout along batch dims.
      dropout_rate: dropout rate
      deterministic: if false, the attention weight is masked randomly using
        dropout, whereas if true, the attention weights are deterministic.
      precision: numerical precision of the computation see `jax.lax.Precision`
        for details.
      kernel_init: initializer for the kernel of the Dense layers.
      bias_init: initializer for the bias of the Dense layers.
      use_bias: bool: whether pointwise QKVO dense transforms use bias.
      attention_fn: dot_product_attention or compatible function. Accepts query,
        key, value, and returns output of shape `[bs, dim1, dim2, ..., dimN,,
        num_heads, value_channels]``
      decode: whether to prepare and use an autoregressive cache.
      normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442).
    """

    num_heads: int
    dtype: Optional[Dtype] = None
    param_dtype: Dtype = jnp.float32
    qk_features: Optional[int] = None
    v_features: Optional[int] = None
    out_features: Optional[int] = None
    broadcast_dropout: bool = True
    dropout_rate: float = 0.0
    deterministic: Optional[bool] = None
    precision: PrecisionLike = None
    kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = default_kernel_init
    bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.zeros_init()
    use_bias: bool = True
    attention_norm: str = "softmax"
    decode: bool = False
    normalize_qk: bool = False
    target_network: str = "default"

    @overload
    def __call__(
        self,
        inputs_q: Array,
        inputs_k: Optional[Array] = None,
        inputs_v: Optional[Array] = None,


Complete the function implementation replacing "raise NotImplementedError".