Skip to content

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

jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:fix-jax-causal-alignment

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

bb6a8cd
Select commit
Loading
Failed to load commit list.
Sign in for the full log view