Repository navigation

Add MultiheadAttention layer - #37

Merged

jessegrabowski merged 4 commits into

mainfrom

causal-self-attention

Jun 15, 2026

Merged

Add MultiheadAttention layer #37
jessegrabowski merged 4 commits into
mainfrom
causal-self-attention

Commits on Jun 15, 2026
- Add scaled dot-product attention and multi-head layers
  jessegrabowski
  committed
  1830cfe
- Add JAX and MLX funcify dispatch for attention
  jessegrabowski
  committed
  3d0f406
- Add jax to test env
  jessegrabowski
  committed
  db733d2
- Add serialization support for MultiheadAttention
  jessegrabowski
  committed
  c59d7d1