diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py
index 1ef71cc..ec2f093 100644
--- a/pytensor_ml/layers/attention.py
+++ b/pytensor_ml/layers/attention.py
@@ -1,7 +1,6 @@
 import numpy as np
 import pytensor.tensor as pt
 
-from pytensor import config
 from pytensor.tensor.variable import TensorVariable
 
 from pytensor_ml.base import Layer, UnaryLayerOp, _resolve_layer_name
@@ -41,10 +40,12 @@ def _sdpa_graph(
     k = _repeat_kv(k, q.type.shape[-3], k.type.shape[-3])
     v = _repeat_kv(v, q.type.shape[-3], v.type.shape[-3])
 
+    input_dtype = q.dtype
+
     if scale is None:
-        scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(config.floatX))
+        scale_t = 1.0 / pt.sqrt(q.shape[-1].astype(input_dtype))
     else:
-        scale_t = pt.as_tensor(scale, dtype=config.floatX)
+        scale_t = pt.as_tensor(scale, dtype=input_dtype)
 
     scores = (q @ k.swapaxes(-1, -2)) * scale_t
 
@@ -54,7 +55,7 @@ def _sdpa_graph(
         # case) sees the whole prefix. Reduces to a plain lower triangle when sq == sk.
         q_idx = pt.arange(sq)[:, None]
         k_idx = pt.arange(sk)[None, :]
-        causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(config.floatX)
+        causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(input_dtype)
         scores = scores + causal
 
     if mask is not None:
diff --git a/tests/test_attention.py b/tests/test_attention.py
index 6fdd1f9..2d97091 100644
--- a/tests/test_attention.py
+++ b/tests/test_attention.py
@@ -50,13 +50,15 @@ def sdpa_np(q, k, v, *, is_causal=False, mask=None, scale=None):
 @pytest.mark.parametrize("is_causal", [False, True], ids=["full", "causal"])
 def test_sdpa_matches_reference(is_causal, scale, rng):
     b, n_head, seq, qk_dim, v_dim = 2, 3, 5, 4, 6
-    q = rng.normal(size=(b, n_head, seq, qk_dim)).astype(floatX)
-    k = rng.normal(size=(b, n_head, seq, qk_dim)).astype(floatX)
-    v = rng.normal(size=(b, n_head, seq, v_dim)).astype(floatX)
+    q = rng.normal(size=(b, n_head, seq, qk_dim)).astype("float32")
+    k = rng.normal(size=(b, n_head, seq, qk_dim)).astype("float32")
+    v = rng.normal(size=(b, n_head, seq, v_dim)).astype("float32")
 
-    out = scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=scale).eval()
+    with pytensor.config.change_flags(floatX="float64"):
+        out = scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=scale).eval()
 
     assert out.shape == (b, n_head, seq, v_dim)
+    assert out.dtype == q.dtype
     np.testing.assert_allclose(out, sdpa_np(q, k, v, is_causal=is_causal, scale=scale), atol=1e-5)
 
 
@@ -74,13 +76,15 @@ def test_sdpa_grouped_query(rng):
 @pytest.mark.parametrize("is_causal", [False, True], ids=["mask_only", "mask_and_causal"])
 def test_sdpa_additive_mask(is_causal, rng):
     b, n_head, seq, dim = 2, 3, 5, 4
-    q = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
-    k = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
-    v = rng.normal(size=(b, n_head, seq, dim)).astype(floatX)
-    mask = np.where(rng.normal(size=(b, n_head, seq, seq)) > 0, 0.0, -1e9).astype(floatX)
+    q = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
+    k = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
+    v = rng.normal(size=(b, n_head, seq, dim)).astype("float32")
+    mask = np.where(rng.normal(size=(b, n_head, seq, seq)) > 0, 0.0, -1e9).astype("float32")
 
-    out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=is_causal).eval()
+    with pytensor.config.change_flags(floatX="float64"):
+        out = scaled_dot_product_attention(q, k, v, mask=mask, is_causal=is_causal).eval()
 
+    assert out.dtype == q.dtype
     np.testing.assert_allclose(out, sdpa_np(q, k, v, mask=mask, is_causal=is_causal), atol=1e-5)