def METHOD_NAME(
    scope: Scope,
    inputs_q: Array,
    inputs_kv: Array,
    bias: Optional[Array] = None,
    qkv_features: Optional[int] = None,
    out_features: Optional[int] = None,
    attn_fn: Callable = softmax_attn,
    dtype=jnp.float32,
):
  if qkv_features is None:
    qkv_features = inputs_q.shape[-1]
  if out_features is None:
    out_features = inputs_q.shape[-1]
  dense = partial(nn.dense, features=qkv_features, bias=False, dtype=dtype)
  query = scope.child(dense, 'query')(inputs_q)
  key = scope.child(dense, 'key')(inputs_kv)
  value = scope.child(dense, 'value')(inputs_kv)
  y = _dot_product_attention(
      scope, query, key, value, bias=bias, attn_fn=attn_fn, dtype=dtype
  )
  return scope.child(nn.dense, 'out')(y, features=out_features, dtype=dtype)