RoPE

1 long-form post on RoPE: machine-learning research by Taha Bouhsine, each built around live, in-browser interactive visualizations.

1 post tagged #rope.

Jun 4, 2026

Q and K Projections in JAX/Flax NNX

A runnable companion to Why Attention Needs Q and K Projections: build scaled dot-product attention with separate query and key projections in Flax NNX, pull the bilinear form B = W_Q W_Kᵀ out of the module, split it into a symmetric metric and an antisymmetric directed part, wire a toy induction head, add RoPE, and measure the low-rank budget and the gauge freedom, all in plain JAX.