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

scaled_dot_product_attention promotes float32 inputs to float64, defeating half-precision graphs

pytensor_ml.layers.scaled_dot_product_attention returns float64 for float32 inputs whenever
pytensor.config.floatX is its default ("float64"). It does so with no float64 constant in the
caller's graph: the promotion comes from config.floatX being used for the softmax scale and the
causal mask inside _sdpa_graph.

Reproducer (self-contained, no downstream project involved)

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

print(pytensor.config.floatX)                 # float64

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))

print(sdpa(q, k, v, is_causal=True).dtype)    # float64  <-- expected float32

Same for AttentionLayer/MultiheadAttention, for scale=None and an explicit scale=, and with or
without is_causal.

Root cause

pytensor_ml/layers/attention.py::_sdpa_graph — three uses of config.floatX where the input dtype is
meant:

scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(config.floatX))     # line 45
scale_t = pt.as_tensor(scale, dtype=config.floatX)             # line 47
causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(config.floatX)   # line 57

q @ k.swapaxes(-1, -2) is float32; multiplying by the float64 scale_t promotes scores, and the
float64 causal mask keeps it there, so softmax(scores) @ v is float64. The Python literals 0.0
and -np.inf are not the problem — pt.where types them to the other branch — only .astype with
config.floatX is.

Impact

Every downstream op computes in float64 with casts at the boundary. Measured in a 30-layer
SmolLM2-135M-shaped decoder consuming this op: 973 of 1743 graph variables become float64 in a model
whose weights are all float32, and decode throughput drops ~10x (3.4 tok/s vs 36.0 tok/s for an
otherwise identical float64 model, i.e. the half-precision path is the slowest way to run it). The
regression is silent: results stay numerically correct, so no value-parity test notices.

Proposed fix

Use the input dtype for both the scale and the causal mask:

compute_dtype = q.dtype
...
scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(compute_dtype))
...
scale_t = pt.as_tensor(scale, dtype=compute_dtype)
...
causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(compute_dtype)

q is the right reference: q/k/v come from the same projections in every realistic call, q @ k^T
already promotes mixed dtypes per NumPy rules, and typing the scale at q.dtype is never wider than
the product. config remains imported elsewhere in the module for MultiheadAttention? — it is used
only on those three lines, so the import can go too.

Verified locally: with the patch, float32/float64 inputs give float32/float64 outputs across
is_causal x scale x mask combinations, with unchanged values.

Same-class defects (separate, not fixed here)

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

Happy to file those separately if you want them.

Start in layers/attention.py at _sdpa_graph and inspect the three config.floatX uses identified in the issue. Reproduce the float32 attention case, update the scale and causal-mask dtype handling as described, then verify float32 and float64 outputs across causal, scale, and mask combinations with unchanged values.