scaled_dot_product_attention

ivy.neural_net_functional.layers.scaled_dot_product_attention(q, k, v, scale, mask=None)[source]

Applies scaled dot product attention to inputs x using optional mask.

Parameters
  • q (array) – The queries [batch_shape,num_queries,feat_dim].

  • k (array) – The keys [batch_shape,num_keys,feat_dim].

  • v (array) – The values [batch_shape,num_keys,feat_dim].

  • scale (float) – The value by which to scale the query-key pairs before softmax.

  • mask (array, optional) – The mask to apply to the query-key values. Default is None. [batch_shape,num_queries,num_keys]

:return The output following application of scaled dot-product attention. [batch_shape,num_queries,feat_dim]


Supported Frameworks:

empty jax_logo empty tf_logo empty pytorch_logo empty mxnet_logo empty numpy_logo empty