# mlx.nn.MultiHeadAttention — MLX 0.0.7 documentation

Skip to main content
`Ctrl`+`K`

# mlx.nn.MultiHeadAttention

class mlx.nn.MultiHeadAttention(dims: int, num_heads: int, query_input_dims: Optional [ int ] = None, key_input_dims: Optional [ int ] = None, value_input_dims: Optional [ int ] = None, value_dims: Optional [ int ] = None, value_output_dims: Optional [ int ] = None, bias: bool = False)

  Implements the scaled dot product attention with multiple heads.

  Given inputs for queries, keys and values the `MultiHeadAttention` produces new values by aggregating information from the input values according to the similarities of the input queries and keys.

  All inputs as well as the output are linearly projected without biases by default.

  `MultiHeadAttention` also takes an optional additive attention mask that should be broadcastable with `(batch, num_heads, # queries, # keys)`. The mask should have `-inf` or very large negative numbers at the positions that should not be attended to.

  Parameters:

    - dims (int) – The model dimensions. This is also the default value for the queries, keys, values, and the output.

    - num_heads (int) – The number of attention heads to use.

    - query_input_dims (int , optional) – The input dimensions of the queries. Default: `dims`.

    - key_input_dims (int , optional) – The input dimensions of the keys. Default: `dims`.

    - value_input_dims (int , optional) – The input dimensions of the values. Default: `key_input_dims`.

    - value_dims (int , optional) – The dimensions of the values after the projection. Default: `dims`.

    - value_output_dims (int , optional) – The dimensions the new values will be projected to. Default: `dims`.

    - bias (bool , optional) – Whether or not to use a bias in the projections. Default: `False`.