Repository navigation

Align the JAX causal mask bottom-right when queries and keys differ - #172

Merged

jessegrabowski merged 2 commits into

pymc-devs:mainfrom

jessegrabowski:fix-jax-causal-alignment

Oct 6, 2026

Merged

Align the JAX causal mask bottom-right when queries and keys differ#172
jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:fix-jax-causal-alignment

Commits on Oct 6, 2026
- Pin the reference causal mask for a query shorter than its keys
  jessegrabowski
  committed
  00e2f2a
- Align the JAX causal mask bottom-right when queries and keys differ
  jessegrabowski
  committed
  bb6a8cd