Fix BaseHMC default mass matrix ignoring requested dtype (GH #8213)

pymc/step_methods/hmc/base_hmc.py

@@ -25,7 +25,6 @@
 from pymc.blocking import DictToArrayBijection, PointType, RaveledVars, StatsType
 from pymc.exceptions import SamplingError
 from pymc.model import Point, modelcontext
-from pymc.pytensorf import floatX
 from pymc.stats.convergence import SamplerWarning, WarningType
 from pymc.step_methods.arraystep import GradientSharedStep
 from pymc.step_methods.compound import StepMethodState
@@ -164,9 +163,16 @@ def __init__(
         self.tune = True
 
         if scaling is None and potential is None:
-            mean = floatX(np.zeros(size))
-            var = floatX(np.ones(size))
-            potential = QuadPotentialDiagAdapt(size, mean, var, 10, rng=self.rng.spawn(1)[0])
+            # Use the same dtype as the logp/dlogp function (which honors the
+            # requested ``dtype``) rather than ``floatX`` unconditionally, so the
+            # default mass matrix always matches the integrator's expected dtype.
+            # See GH #8213.
+            dtype = self._logp_dlogp_func.dtype
+            mean = np.zeros(size, dtype=dtype)
+            var = np.ones(size, dtype=dtype)
+            potential = QuadPotentialDiagAdapt(
+                size, mean, var, 10, dtype=dtype, rng=self.rng.spawn(1)[0]
+            )
 
         if isinstance(scaling, dict):
             point = Point(scaling, model=self._model)

tests/step_methods/hmc/test_hmc.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 import numpy.testing as npt
+import pytensor
 import pytest
 
 import pymc as pm
@@ -88,3 +89,35 @@ def test_nuts_tuning():
     ss_tuned = idata.warmup_sample_stats["step_size"][0, -1]
     ss_posterior = idata.sample_stats["step_size"][0, :]
     np.testing.assert_array_equal(ss_posterior, ss_tuned)
+
+
+def test_default_potential_honors_dtype():
+    """The default mass matrix must use the sampler's dtype, not ``floatX``.
+
+    ``BaseHMC`` forwards ``dtype`` to the logp/dlogp function but used to build
+    the default ``QuadPotentialDiagAdapt`` with ``floatX`` unconditionally. When
+    the requested ``dtype`` differed from ``pytensor.config.floatX`` the potential
+    and the logp function ended up with mismatched dtypes and
+    ``CpuLeapfrogIntegrator`` raised ``ValueError: dtypes of potential ... and
+    logp function ... don't match``. See GH #8213.
+    """
+
+    class HMC(BaseHMC):
+        def _hamiltonian_step(self, *args, **kwargs):
+            pass
+
+    # Build a float64 model, then construct the step method while the global
+    # floatX is float32. The requested dtype (float64) must win for the
+    # potential, matching the logp function, so no dtype mismatch is raised.
+    with pm.Model() as model:
+        pm.Normal("x", 0.0, 1.0)
+
+    assert all(v.dtype == "float64" for v in model.value_vars)
+
+    with pytensor.config.change_flags(floatX="float32"):
+        step = HMC(vars=model.value_vars, model=model, dtype="float64")
+
+    assert step.potential.dtype == np.dtype("float64")
+    assert step.potential._var.dtype == np.dtype("float64")
+    # The integrator's own check must agree (it raises on a mismatch on init).
+    assert step.integrator._potential.dtype == step.integrator._dtype