PyTorch
API: torch.nn.modules.transformer.TransformerDecoderLayer.norm3
Strategy: Direct Mapping
JAX (Core)
API: optax.tree.norm
Strategy: Direct Mapping
Flax NNX
API: optax.tree.norm
Strategy: Direct Mapping
PaxML / Praxis
API: optax.tree.norm
Strategy: Direct Mapping