BUG: BaseHMC doesn't pass `dtype` to ‎`QuadPotentialDiagAdapt`

Start in pymc/step_methods/hmc/base_hmc.py at the QuadPotentialDiagAdapt construction, then read pymc/step_methods/hmc/integration.py where the potential and logp function dtypes are compared. Inspect the existing HMC and quad-potential tests before making the change. Done means an explicitly supplied dtype reaches the potential and avoids the reported mismatch without breaking default behavior.

Describe the issue:

By inspection (and in a test of my code), BaseHMC accepts a dtype parameter, but in some cases creates a ‎QuadPotentialDiagAdapt instance without passing the dtype to it. In my case, this appears to be the source of a ValueError.

I'd hoped to pass dtypes explicitly to my PYMC objects to keep my code maximally flexible, e.g. eventually allowing models to be used in parallel threads in the same process, but at different levels of precision.

I'm aware of the pytensor.config.floatX option that should fix this, but I thought it potentially useful to report the use of local state in PYMC that seems likely to fix this without the need for less flexible global configuration.

Attaching a partial traceback below with just the PYMC-related portion.

Reproduceable code example:

[Model constructed using explicit `dtype` params]

.NUTS(model=model, dtype=dtype_str)

Error message:

step=pm.NUTS(model=model, dtype=dtype_str), 
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/app.venv/lib/python3.12/site-packages/pymc/step_methods/hmc/nuts.py:198: in __init__
    super().__init__(vars, **kwargs)
/app.venv/lib/python3.12/site-packages/pymc/step_methods/hmc/base_hmc.py:183: in __init__
    self.integrator = integration.CpuLeapfrogIntegrator(self.potential, self._logp_dlogp_func)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <pymc.step_methods.hmc.integration.CpuLeapfrogIntegrator object at 0xffff88d9b590>, potential = <pymc.step_methods.hmc.quadpotential.QuadPotentialDiagAdapt object at 0xffff88878ce0>
logp_dlogp_func = <pymc.model.core.ValueGradFunction object at 0xffff89c55c70>

    def __init__(self, potential: QuadPotential, logp_dlogp_func):
        """Leapfrog integrator using CPU."""
        self._potential = potential
        # Sidestep logp_dlogp_function.__call__
        pytensor_function = logp_dlogp_func._pytensor_function
        # Create some wrappers for backwards compatibility during transition
        # When raveled_inputs=False is forbidden, func = pytensor_function
        if logp_dlogp_func._raveled_inputs:
    
            def func(q, _):
                return pytensor_function(q)
    
        else:
    
            def func(q, point_map_info):
                unraveled_q = DictToArrayBijection.rmap(RaveledVars(q, point_map_info)).values()
                return pytensor_function(*unraveled_q)
    
        self._logp_dlogp_func = func
        self._dtype = logp_dlogp_func.dtype
        if self._potential.dtype != self._dtype:
>           raise ValueError(
                f"dtypes of potential ({self._potential.dtype}) and logp function ({self._dtype})"
                "don't match."
            )
E           ValueError: dtypes of potential (float64) and logp function (float32)don't match.

/app.venv/lib/python3.12/site-packages/pymc/step_methods/hmc/integration.py:61: ValueError

PyMC version information:

pymc: 5.23.0 (linked PYMC source is from master though, where this seems to still exist)
pytensor: 2.31.7
Installed in a Debian Docker container using uv. Host OS is MacOS Sequoia.

Context for the issue:

Reduces flexibility of PYMC-based code using this code path. Forces clients to always use the same precision during a run or to use workarounds to use models with different precision in the same process.