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: ValueErrorPyMC 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.
Describe the issue:
By inspection (and in a test of my code),
BaseHMCaccepts adtypeparameter, but in some cases creates a QuadPotentialDiagAdaptinstance without passing thedtypeto 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:
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: ValueErrorPyMC 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.