diff --git a/docs/source/api/model/optimization.rst b/docs/source/api/model/optimization.rst index ffd23cf89f..4a62efb8f4 100644 --- a/docs/source/api/model/optimization.rst +++ b/docs/source/api/model/optimization.rst @@ -6,3 +6,5 @@ Model Optimization freeze_dims_and_data freeze_model + model_to_float32 + model_to_float64 diff --git a/pymc/model/transform/__init__.py b/pymc/model/transform/__init__.py index 1dea85a61b..fa9f5c3be7 100644 --- a/pymc/model/transform/__init__.py +++ b/pymc/model/transform/__init__.py @@ -28,7 +28,12 @@ extract_deterministics, insert_deterministics, ) -from pymc.model.transform.optimization import freeze_dims_and_data, freeze_model +from pymc.model.transform.optimization import ( + freeze_dims_and_data, + freeze_model, + model_to_float32, + model_to_float64, +) __all__ = ( "change_value_transforms", @@ -37,6 +42,8 @@ "freeze_dims_and_data", "freeze_model", "insert_deterministics", + "model_to_float32", + "model_to_float64", "observe", "prune_vars_detached_from_observed", "remove_minibatched_nodes", diff --git a/pymc/model/transform/optimization.py b/pymc/model/transform/optimization.py index 87c6bba29f..cdffdf9471 100644 --- a/pymc/model/transform/optimization.py +++ b/pymc/model/transform/optimization.py @@ -11,18 +11,32 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy + from collections.abc import Sequence from typing import cast import numpy as np +import pytensor +import pytensor.tensor as pt from pytensor.compile import SharedVariable +from pytensor.compile.builders import OpFromGraph, construct_nominal_fgraph from pytensor.graph import Constant, FunctionGraph, Variable from pytensor.graph.replace import clone_replace -from pytensor.graph.traversal import ancestors +from pytensor.graph.traversal import ancestors, io_toposort +from pytensor.scalar import Cast +from pytensor.tensor.elemwise import Elemwise +from pytensor.tensor.type import TensorType +from pymc.logprob.transforms import Transform from pymc.model.core import FrozenModel, Model -from pymc.model.fgraph import ModelFreeRV, fgraph_from_model, model_from_fgraph +from pymc.model.fgraph import ( + ModelFreeRV, + ModelValuedVar, + fgraph_from_model, + model_from_fgraph, +) def _constant_from_shared(shared: SharedVariable) -> Constant: @@ -229,4 +243,220 @@ def freeze_model(model: Model) -> FrozenModel: return cast(FrozenModel, frozen_model) -__all__ = ("freeze_dims_and_data", "freeze_model") +def _is_dtype(dtype, ref_dtype: str) -> bool: + """Whether `dtype` (a dtype-like or alias such as "float") normalizes to `ref_dtype`.""" + try: + return dtype is not None and np.dtype(dtype).name == ref_dtype + except TypeError: + return False + + +def _cast_root(var: Variable, from_dtype: str, to_dtype: str) -> Variable: + """Return a `to_dtype` clone of a root variable (constant, shared or input).""" + if getattr(var.type, "dtype", None) != from_dtype: + return var + if isinstance(var, Constant): + return pt.constant(var.data.astype(to_dtype), name=var.name) + if isinstance(var, SharedVariable): + return pytensor.shared( + var.get_value(borrow=False).astype(to_dtype), name=var.name, shape=var.type.shape + ) + return var.type.clone(dtype=to_dtype)(name=var.name) + + +def _restore_static_shape(new: Variable, old: Variable) -> Variable: + if isinstance(new.type, TensorType) and new.type.shape != old.type.shape: + new = pt.specify_shape(new, old.type.shape) + new.name = old.name + return new + + +class _CastedTransform(Transform): + """Wrap a transform whose graphs produce a different float dtype. + + The wrapped transform's graphs are converted with `_cast_graph_floats` — so + constants of the old dtype embedded in the transform (not reachable from the + model graph) are cast too — and the outputs are cast as a last resort. + """ + + def __init__(self, transform: Transform, from_dtype: str, to_dtype: str): + self.transform = transform + self.from_dtype = from_dtype + self.to_dtype = to_dtype + # Keep the name: value variable names derive from it + self.name = transform.name + + def _converted(self, out: Variable) -> Variable: + (out,), _ = _cast_graph_floats([out], self.from_dtype, self.to_dtype) + return pt.cast(out, self.to_dtype) + + def forward(self, value, *inputs): + return self._converted(self.transform.forward(value, *inputs)) + + def backward(self, value, *inputs): + return self._converted(self.transform.backward(value, *inputs)) + + def log_jac_det(self, value, *inputs): + return self._converted(self.transform.log_jac_det(value, *inputs)) + + +def _transform_keeps_dtype(transform: Transform, rv: Variable, value: Variable, dtype: str) -> bool: + """Whether the transform's forward/backward graphs on `rv`/`value` stay in `dtype`. + + Probed under ``floatX=dtype``, the setting the converted model is meant to be + compiled under, so only transforms that embed foreign-dtype constants get wrapped. + """ + try: + with pytensor.config.change_flags(floatX=dtype): + return ( + transform.forward(rv, *rv.owner.inputs).type.dtype == dtype + and transform.backward(value, *rv.owner.inputs).type.dtype == dtype + ) + except Exception: + return False + + +def _cast_graph_floats( + outputs: Sequence[Variable], from_dtype: str, to_dtype: str +) -> tuple[list[Variable], dict[Variable, Variable]]: + """Clone the graph of `outputs`, casting every `from_dtype` variable to `to_dtype`. + + Returns the converted outputs and a memo mapping old to new variables. + """ + memo: dict[Variable, Variable] = {} + + def mapped(var): + if var not in memo: + memo[var] = _cast_root(var, from_dtype, to_dtype) + return memo[var] + + for node in io_toposort([], outputs): + op, new_inputs = node.op, [mapped(var) for var in node.inputs] + if ( + isinstance(op, Elemwise) + and isinstance(op.scalar_op, Cast) + and op.scalar_op.o_type.dtype == from_dtype + ): + # Redirect explicit casts (e.g. `x.astype("float64")`) + new_outputs = [pt.cast(new_inputs[0], to_dtype)] + elif isinstance(op, OpFromGraph): + # Convert the inner graph of e.g. SymbolicRandomVariables recursively. + # Static shapes frozen in the inner graph cannot be re-inferred from the + # inner inputs when nodes are rebuilt, so they are restored explicitly. + inner_outs, inner_memo = _cast_graph_floats(op.inner_outputs, from_dtype, to_dtype) + inner_outs = [ + _restore_static_shape(new, old) + for new, old in zip(inner_outs, op.inner_outputs, strict=True) + ] + inner_ins = [ + inner_memo.get(i, _cast_root(i, from_dtype, to_dtype)) for i in op.inner_inputs + ] + new_op = copy.copy(op) + new_op.fgraph = construct_nominal_fgraph(inner_ins, inner_outs).freeze() + new_op.input_types = [i.type for i in inner_ins] + new_op.output_types = [o.type for o in inner_outs] + # Drop gradient caches computed for the old inner graph + new_op._lop_op_cache = {} + new_op._rop_op_cache = None + new_op._frozen_lop = None + new_op._frozen_rop = None + new_outputs = new_op.make_node(*new_inputs).outputs + elif isinstance(op, ModelValuedVar) and op.transform is not None: + # Transform objects travel with the op and may embed constants of the old + # dtype in the value-space graphs (logp, initial point); wrap them if so. + rv_new, value_new = new_inputs + if not _transform_keeps_dtype(op.transform, rv_new, value_new, to_dtype): + op = copy.copy(op) + op.transform = _CastedTransform(op.transform, from_dtype, to_dtype) + new_outputs = op.make_node(*new_inputs).outputs + elif _is_dtype(getattr(op, "dtype", None), from_dtype): + # Ops with a fixed output dtype: RandomVariables, reductions, ARange, ... + new_op = copy.copy(op) + new_op.dtype = to_dtype + new_outputs = new_op.make_node(*new_inputs).outputs + else: + new_outputs = op.make_node(*new_inputs).outputs + for old, new in zip(node.outputs, new_outputs, strict=True): + new.name = old.name + memo[old] = new + + return [mapped(out) for out in outputs], memo + + +def _cast_model_floats(model: Model, from_dtype: str, to_dtype: str) -> Model: + initial_values = _extract_initial_values(model) + saved_initial_values = dict(model.rvs_to_initial_values) + try: + for rv in model.rvs_to_initial_values: + model.rvs_to_initial_values[rv] = None + fg, _ = fgraph_from_model(model) + finally: + model.rvs_to_initial_values.update(saved_initial_values) + + new_outputs, memo = _cast_graph_floats(fg.outputs, from_dtype, to_dtype) + new_fg = FunctionGraph(outputs=new_outputs, clone=False) + new_fg._coords = fg._coords # type: ignore[attr-defined] + new_fg._dim_lengths = { # type: ignore[attr-defined] + dim: memo.get(length, length) + for dim, length in fg._dim_lengths.items() # type: ignore[attr-defined] + } + + new_model = model_from_fgraph(new_fg, mutate_fgraph=True) + for name, initval in initial_values.items(): + if isinstance(initval, np.ndarray) and initval.dtype.kind == "f": + initval = initval.astype(to_dtype) + new_model.set_initval(new_model[name], initval) + return new_model + + +def model_to_float32(model: Model) -> Model: + """Recreate a Model with all float64 variables and data cast to float32. + + Every float64 variable is converted: data (constants and `pm.Data`), free and + observed RVs (including the inner graphs of symbolic RVs like `ZeroSumNormal`), + value variables, Deterministics and Potentials. Integer, boolean and RNG + variables are unaffected. Explicit `.astype("float64")` casts are redirected + to float32. + + This can speed up sampling at the cost of precision — most on GPUs and for + compute-bound models; on CPU backends gains depend on how memory- and + BLAS-bound the model's logp is. + + Compile and sample under ``floatX="float32"``, otherwise constants introduced + when building logp graphs will upcast intermediate computations back to float64: + + .. code-block:: python + + import pymc as pm + import pytensor + + from pymc.model.transform.optimization import model_to_float32 + + with pm.Model() as m: + x = pm.Data("x", [0.0, 1.0, 2.0]) + beta = pm.Normal("beta") + pm.Normal("y", mu=beta * x, sigma=1.0, observed=[1.0, 2.0, 3.0]) + + with pytensor.config.change_flags(floatX="float32"): + with model_to_float32(m): + idata = pm.sample() + + Notes + ----- + ``pm.set_data`` on the new model expects float32 arrays. + + Constant and strategy-string initial values are preserved (arrays are cast); + symbolic initial values are not supported. + """ + return _cast_model_floats(model, "float64", "float32") + + +def model_to_float64(model: Model) -> Model: + """Recreate a Model with all float32 variables and data cast to float64. + + The inverse of :func:`model_to_float32`. See its docstring for details. + """ + return _cast_model_floats(model, "float32", "float64") + + +__all__ = ("freeze_dims_and_data", "freeze_model", "model_to_float32", "model_to_float64") diff --git a/tests/model/transform/test_optimization.py b/tests/model/transform/test_optimization.py index 326e9bb7c3..98433d7b3e 100644 --- a/tests/model/transform/test_optimization.py +++ b/tests/model/transform/test_optimization.py @@ -12,17 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. import numpy as np +import pytensor import pytest from pytensor.compile import SharedVariable from pytensor.graph import Constant +import pymc as pm + from pymc import Deterministic, do from pymc.data import Data from pymc.distributions import HalfNormal, Normal from pymc.exceptions import NotConstantValueError from pymc.model import Model -from pymc.model.transform.optimization import freeze_dims_and_data +from pymc.model.transform.optimization import ( + freeze_dims_and_data, + model_to_float32, + model_to_float64, +) from pymc.pytensorf import constant_fold @@ -179,3 +186,124 @@ def test_freeze_dims_and_data_partially_observed_rv(): frozen_y = freeze_dims_and_data(model)["y"] assert constant_fold([frozen_y.shape]) == (3,) + + +class TestModelToFloat32: + @staticmethod + def _mixed_model(): + rng = np.random.default_rng(4) + x_data = rng.normal(size=10) + with Model(coords={"g": range(3)}) as m: + x = Data("x", x_data) + idx = Data("idx", np.arange(10)) + beta = Normal("beta") + sigma = HalfNormal("sigma") + z = pm.ZeroSumNormal("z", dims="g") + det = Deterministic("det", beta * x + z.mean()) + Normal("y", mu=det[idx], sigma=sigma, observed=x_data * 2) + return m + + def test_dtypes_converted(self): + m = self._mixed_model() + m32 = model_to_float32(m) + + for name in ("x", "beta", "sigma", "z", "det", "y"): + assert m32[name].type.dtype == "float32", name + assert m32["idx"].type.dtype == m["idx"].type.dtype + for rv in m32.free_RVs + m32.observed_RVs: + assert m32.rvs_to_values[rv].type.dtype == "float32" + # Static shapes and transforms are preserved + for name in ("x", "beta", "sigma", "z", "det", "y"): + assert m32[name].type.shape == m[name].type.shape, name + assert type(m32.rvs_to_transforms[m32["z"]]) is type(m.rvs_to_transforms[m["z"]]) + with Model() as m_static: + pm.ZeroSumNormal("z", shape=(3,)) + assert model_to_float32(m_static)["z"].type.shape == (3,) + + def test_logp_and_draws(self): + m = self._mixed_model() + m32 = model_to_float32(m) + + ip64 = m.initial_point() + logp64 = m.compile_logp()(ip64) + with pytensor.config.change_flags(floatX="float32"): + ip32 = m32.initial_point() + assert all(v.dtype == "float32" for v in ip32.values()) + logp32 = m32.compile_logp()(ip32) + dlogp32 = m32.compile_dlogp()(ip32) + np.testing.assert_allclose(logp32, logp64, rtol=1e-5) + assert np.asarray(dlogp32).dtype == "float32" + assert pm.draw(m32["z"], random_seed=1).dtype == "float32" + + def test_round_trip(self): + m = self._mixed_model() + m64 = model_to_float64(model_to_float32(m)) + for name in ("x", "beta", "sigma", "z", "det", "y"): + assert m64[name].type.dtype == "float64", name + np.testing.assert_allclose( + m64.compile_logp()(m64.initial_point()), + m.compile_logp()(m.initial_point()), + rtol=1e-5, + ) + + def test_explicit_cast_redirected(self): + with Model() as m: + x = Data("x", np.arange(5)) # int64 + Normal("y", mu=x.astype("float64"), observed=np.zeros(5)) + m32 = model_to_float32(m) + assert m32["y"].type.dtype == "float32" + assert m32["y"].owner.inputs[3].type.dtype == "float32" + + def test_preserves_initvals(self): + with Model() as m: + sigma = HalfNormal("sigma", initval=np.array(5.0)) + beta = Normal("beta", initval="prior") + m32 = model_to_float32(m) + initval = m32.rvs_to_initial_values[m32["sigma"]] + assert initval.dtype == "float32" and initval == 5.0 + assert m32.rvs_to_initial_values[m32["beta"]] == "prior" + + def test_sample_smoke(self): + m32 = model_to_float32(self._mixed_model()) + with pytensor.config.change_flags(floatX="float32"): + with m32: + idata = pm.sample( + draws=10, + tune=10, + chains=1, + progressbar=False, + random_seed=1, + compute_convergence_checks=False, + ) + # The sampler may store draws as float64; just check sampling worked + assert np.isfinite(idata.posterior["beta"]).all() + + def test_transform_with_foreign_dtype_constants(self): + # Transforms travel with the model as objects and may bake float64 constants + # into value-space graphs; model_to_float32 must keep those graphs float32. + import pytensor.tensor as pt + + from pymc.logprob.transforms import Transform + + class ScaledTransform(Transform): + name = "scaled" + scale = np.array(2.0, dtype="float64") # baked float64 constant + + def forward(self, value, *inputs): + return value * self.scale + + def backward(self, value, *inputs): + return value / self.scale + + def log_jac_det(self, value, *inputs): + return -pt.log(self.scale) * pt.ones_like(value) + + with Model() as m: + Normal("x", 0, 1, default_transform=ScaledTransform()) + + m32 = model_to_float32(m) + with pytensor.config.change_flags(floatX="float32"): + ip32 = m32.initial_point() + assert all(v.dtype == "float32" for v in ip32.values()) + logp32 = m32.compile_logp()(ip32) + np.testing.assert_allclose(logp32, m.compile_logp()(m.initial_point()), rtol=1e-5)