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

[Mohit-Ak]

Description

BaseHMC.__init__ accepts a dtype argument and forwards it to the parent so the compiled logp/dlogp function is built with that dtype. However, when it constructs the default mass matrix it ignored that dtype and used floatX unconditionally:

mean = floatX(np.zeros(size))
var = floatX(np.ones(size))
potential = QuadPotentialDiagAdapt(size, mean, var, 10, rng=...)

So whenever the requested dtype differs from pytensor.config.floatX, the potential and the logp function end up with different dtypes, and CpuLeapfrogIntegrator.__init__ bails out:

ValueError: dtypes of potential (float32) and logp function (float64) don't match.

That traceback matches the one in #8213 exactly (the crash originates from the QuadPotentialDiagAdapt(...) line the reporter linked to).

The fix builds the default mass matrix using the dtype the logp/dlogp function actually resolved to (self._logp_dlogp_func.dtype) — which is the same value the integrator later compares against — instead of floatX. That keeps the two in sync in every case, including the common one where dtype is None (both fall back to floatX, so behavior is unchanged). The now-unused floatX import is dropped.

The scaling-based quad_potential(...) path already derives its dtype from the user-supplied array, so it isn't affected; this was the only default-construction site that hard-coded floatX.

How was it tested

Added test_default_potential_honors_dtype in tests/step_methods/hmc/test_hmc.py. It builds a float64 model, then (inside pytensor.config.change_flags(floatX="float32")) constructs the step with dtype="float64" and asserts the default potential comes out float64 and matches the integrator's expected dtype. On main this test fails with the don't match ValueError; with the fix it passes.

Local run:

• pytest tests/step_methods/hmc/ → 67 passed, 6 skipped (matplotlib-optional tests run once installed).
• ruff==0.11.13 ruff check and ruff format --check both clean on the changed files.

Related Issue

• Closes BUG: BaseHMC doesn't pass dtype to ‎QuadPotentialDiagAdapt #8213

Checklist

• Checked that the pre-commit linting/style checks pass (ruff 0.11.13)
• Included tests that prove the fix is effective
• Added necessary documentation (inline comment explaining the dtype choice; no public API change)

Type of change

• Bug fix

[welcome]

]
💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #33386953 | 📁 Comparing 5c3adc5 against latest (c032457)

  🔍 Preview build  

159 files changed · ± 159 modified

± Modified

• glossary.html
• api/math.html
• contributing/implementing_distribution.html
• api/distributions/custom.html
• api/generated/pymc.Data.html
• api/generated/pymc.adadelta.html
• api/generated/pymc.backends.zarr.ZarrChain.html
• api/generated/pymc.draw.html
• api/generated/pymc.init_nuts.html
• api/generated/pymc.ode.DifferentialEquation.html
• and 149 more...