MultiheadAttention

Auto-generated from flax_nnx_code_defs

PyTorchKerasTensorFlowApple MLXFlax NNX

PyTorch

API: torch.nn.MultiheadAttention

Strategy: Direct Mapping

Official Docs ↗

Keras

API: keras.layers.MultiHeadAttention

Strategy: Plugin (repack_attn_keras)

TensorFlow

API: keras.layers.MultiHeadAttention

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: mlx.nn.layers.transformer.MultiHeadAttention

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.MultiHeadAttention

Strategy: Plugin (repack_attn_flax)