AttentionLayer

Dot-product attention layer, a.k.a. Luong-style attention.

Abstract Signature:

AttentionLayer(use_scale: bool = False, score_mode: str = dot, dropout: float = 0.0, seed: int)

PyTorchJAX (Core)KerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: torch.nn.MultiheadAttention

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: —

Strategy: Custom / Partial

Keras

API: keras.layers.Attention

Strategy: Direct Mapping

Official Docs ↗

TensorFlow

API: tf.keras.layers.Attention

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: —

Strategy: Custom / Partial

Flax NNX

API: flax.linen.attention.dot_product_attention

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: —

Strategy: Custom / Partial