AttentionProjection

Layer that computes multi-head projection for dot-product attention.

Abstract Signature:

AttentionProjection(input_dim: int, num_heads: int, dim_per_head: int, is_output: bool = False)

PyTorchFlax NNXPaxML / Praxis

PyTorch

API: torch.nn.Linear

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.Linear

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: praxis.layers.AttentionProjection

Strategy: Direct Mapping

Official Docs ↗