Skip to content

Build attention's scale and causal mask at the queries' dtype - #173

Merged
jessegrabowski merged 1 commit into
pymc-devs:mainfrom
jessegrabowski:fix-attention-dtype
Oct 6, 2026
Merged

jessegrabowski merged 1 commit into
pymc-devs:mainfrom
jessegrabowski:fix-attention-dtype

Conversation

@jessegrabowski

@jessegrabowski jessegrabowski commented Oct 6, 2026 •

Copy link
Copy Markdown
Member

scaled_dot_product_attention now builds its scale and causal mask at the queries' dtype instead of floatX, so float32 attention stays float32 when floatX is float64.

Closes #164


📚 Documentation preview 📚: https://pytensor-ml--173.org.readthedocs.build/en/173/

@jessegrabowski
jessegrabowski merged commit 0784f86 into pymc-devs:main Oct 6, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

scaled_dot_product_attention promotes float32 inputs to float64 (uses config.floatX for the scale and causal mask)

1 participant