ALiBi

Attention with Linear Biases (ALiBi) layer.

PyTorchKerasApple MLXFlax NNX

PyTorch

API: —

Strategy: Custom / Partial

Keras

API: —

Strategy: Custom / Partial

Apple MLX

API: mlx.nn.ALiBi

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.ALiBi

Strategy: Direct Mapping

Official Docs ↗