Fix SDPA output dtype following config.floatX instead of the input dtype

[cetagostini]

Fix: SDPA output dtype follows the input dtype, not config.floatX

Closes #164.

Root cause

_sdpa_graph used pytensor.config.floatX for three tensors: the softmax scale (both the default 1/sqrt(d_k) path and the explicit scale= path) and the causal mask. When config.floatX is "float64" (the default), a float32 (q @ k^T) is promoted to float64 by the scale multiplication, and the float64 causal mask keeps it there.

Fix

Replace config.floatX with q.dtype (named input_dtype) at the three call sites. q is the right reference: q/k/v come from the same projections in every realistic call, and q @ k^T already promotes mixed dtypes per NumPy rules.

Minimal reproducer

import pytensor, pytensor.tensor as pt
from pytensor_ml.layers import scaled_dot_product_attention as sdpa

q = pt.tensor4("q", dtype="float32", shape=(1, 4, 1, 64))
k = pt.tensor4("k", dtype="float32", shape=(1, 2, 8, 64))
v = pt.tensor4("v", dtype="float32", shape=(1, 2, 8, 64))

# Before: float64   After: float32
print(sdpa(q, k, v, is_causal=True).dtype)

Value delta

The only numerical change is that the softmax scale is now computed in float32 instead of float64. On a typical d_k = 64, the float32 scale 1/sqrt(64) differs from the float64 value by < 2e-7. The existing SDPA reference and additive-mask tests now also assert that the output dtype matches the input dtype.

Tests

Dtype assertions are folded into the existing test_sdpa_matches_reference and test_sdpa_additive_mask tests, rather than adding standalone tests. Their six existing cases use float32 inputs with config.floatX="float64", covering default/custom scales, causal/non-causal attention, and additive masks. The existing numerical comparisons are retained.

• pytest tests/test_attention.py -q: 25 passed.
• Ruff lint and formatting checks pass for both changed files.
• An in-memory negative control restoring the original config.floatX behavior fails all six cases at the dtype assertion; the fixed implementation passes all six. The control comparison used PYTENSOR_FLAGS=cxx= to avoid a local C-linker configuration issue; the full 25-test run above used the normal configuration.

Pre-existing quirks NOT changed by this PR

1. float16 inputs produce float32 output — the Python literal 1.0 in 1.0 / pt.sqrt(...) is typed float32 by PyTensor, so the scale is float32 even when q is float16. This is pre-existing and out of scope.
2. A float64 additive mask promotes a float32 graph — if the caller supplies a float64 mask, NumPy promotion widens the scores. Also pre-existing and out of scope.

Same-class defects (follow-ups, not fixed here)

• layers/dropout.py:87 — mask.astype(config.floatX)
• layers/conv.py:1314-1317 — pt.cast(extent, config.floatX) in the interpolation path

---

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

[jessegrabowski]

Suggested change

	compute_dtype = q.dtype
	input_dtype = q.dtype

A bit more clear

[jessegrabowski]

This is way too many tests, probably you can just add a check to the existing ones

[cetagostini]

Both points addressed in 7e8f24d:

• renamed to input_dtype (clearer than compute_dtype).
• the four new tests are gone; the dtype check now lives in the existing tests. test_sdpa_matches_reference and test_sdpa_additive_mask build float32 inputs under change_flags(floatX="float64") and assert out.dtype == q.dtype, which covers both the is_causal and scale/mask paths with one line each — net tests/test_attention.py is now +13/−9 instead of +56.

The check still catches the bug it is for: reverting attention.py to 919d8fd while keeping the revised tests fails all six parametrisations (assert dtype('float64') == dtype('float32')), and the full file is 25 passed with the fix. CI is green on this head (unittest groups, pre-commit, mypy, docs preview).

Downstream (pymc-labs/pytensor.cpp) this unblocks removing our boundary-cast workaround; it currently pins this commit because no release post-dates v0.2.2. No rush implied — just flagging that a release containing it is what lets us unpin.