pymc.model.transform.model_to_float32#

pymc.model.transform.model_to_float32(model)[source]#

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:

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.