Add model_to_float32 / model_to_float64 model transforms

[velochy]

Add model_to_float32 / model_to_float64 model transforms

With floatX="float32", pytensor still permits mixed precision — unlike JAX, which refuses to upcast without x64 enabled. Any float64 constant, pm.Data, or RV already baked into the model graph silently upcasts everything downstream, so models "built float64" never actually sample in single precision — in fact pm.sample under float32 flags fails outright on such models, because the init jitter builds float32 RVs against the float64 graph. This PR makes precision a one-call property of the model instead of something threaded through every input:

from pymc.model.transform.optimization import model_to_float32

with pytensor.config.change_flags(floatX="float32"):
    with model_to_float32(model):
        idata = pm.sample()

model_to_float32(model) recreates the model via the fgraph round-trip with every float64 variable cast to float32: data (constants and pm.Data shareds), RandomVariables and other ops with a baked output dtype (dtype aliases like AllocEmpty{dtype='float'} are normalized), the inner graphs of SymbolicRandomVariables (rebuilt recursively, with frozen static shapes restored via specify_shape), value variables, and explicit .astype("float64") casts. Value transforms are probed under the target floatX, and any transform that embeds foreign-dtype constants (e.g. precomputed float64 basis matrices) is wrapped so its forward/backward graphs are converted too. Integer, boolean, RNG and dim-length variables are untouched; constant and strategy-string initial values are preserved. model_to_float64 is the inverse, sharing the same machinery.

Validated end-to-end on a production survey-modeling pipeline sampled with nutpie: logp matches float64 to ~1e-8 relative, posteriors agree within MC error, and the optimized logp+dlogp graphs come out fully float32. Speed is model- and backend-dependent: on a CPU/numba backend these (gather/indexing-bound) models gain only a few percent per gradient, measured with interleaved A/B timing; the larger wins are expected on GPU backends and for compute-bound models.

Limitations: the converted model should be compiled under floatX="float32", since logp graphs are built lazily and Python-float constants introduced there follow floatX (documented in the docstring). pm.set_data on the converted model expects float32 arrays. The OpFromGraph rebuild currently reaches into private attrs (fgraph, input_types, output_types, gradient caches); happy to move that behind a proper pytensor helper as a follow-up if preferred.

Tests cover dtype conversion across RV kinds (incl. ZeroSumNormal, Deterministics, observed data), logp/dlogp equality vs float64, sampling smoke incl. nutpie, static-shape and transform preservation, foreign-dtype transform wrapping, initval preservation, explicit-cast redirection, and the float32→float64 round trip. Docs: added to docs/source/api/model/optimization.rst with a usage example.

[velochy]

To be very clear - I do not think this is a great idea as it stands. It's a proof of concept to what @jessegrabowski proposed off-hand in discord.

For one, you probably want this as a setting for freeze_dims_and_data or as a standard pre-processing step before sampling all models (to unify inputs to either F32 or F64 world) instead of just stand-alone userspace functions.

Also, it is worth pointing out that the functions cannot actually convert all inputs - in quite a few cases, supposedly. Quoting my AI friend:

model_to_float32 can only convert float64 that is already in the model's fgraph at transform time; anything materialized later — during lazy construction of logp/dlogp/initial-point/PPC graphs — escapes it, so manual floatX-hygiene is required wherever a Python callable or stateful object injects constants at build time: CustomDist-style logp/logcdf/random/support_point closures over numpy arrays (the SIP weights case), any explicit numpy dtype (np.array, np.float64, np.linalg outputs) in lazily-executed code (floatX governs Python floats, never explicit numpy dtypes), and graph-factory objects without an interceptable interface (value transforms are auto-wrapped via their narrow forward/backward/log_jac_det interface, but e.g. GP conditional calls post-transform stitch onto the old graph); additionally, runtime feeds must match the new dtype (pm.set_data needs f32 arrays; old shared references point to dead containers) and symbolic initvals are rejected and must be re-specified — the rule: keep float constants either inside the model graph or created via floatX-respecting paths (Python floats, pm.floatX), never as raw f64 numpy baked into callables.

But this should serve as a proof-of-concept in terms of what it would take, as a discussion starter.

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #34109000 | 📁 Comparing 2d88df4 against latest (da8fc47)

  🔍 Preview build  

9 files changed · + 2 added · ± 7 modified

+ Added

• api/model/generated/pymc.model.transform.model_to_float32.html
• api/model/generated/pymc.model.transform.model_to_float64.html

± Modified

• genindex.html
• glossary.html
• api/model.html
• api/model/fgraph.html
• api/model/optimization.html
• api/model/generated/pymc.model.transform.freeze_model.html
• _modules/pymc/model/transform/optimization.html

[codecov]

Codecov Report

❌ Patch coverage is 96.39640% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.90%. Comparing base (da8fc47) to head (2d88df4).

Files with missing lines	Patch %	Lines
pymc/model/transform/optimization.py	96.36%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8402      +/-   ##
==========================================
+ Coverage   91.87%   91.90%   +0.02%     
==========================================
  Files         128      128              
  Lines       21256    21363     +107     
==========================================
+ Hits        19530    19633     +103     
- Misses       1726     1730       +4     

Files with missing lines	Coverage Δ
pymc/model/transform/__init__.py	100.00% <100.00%> (ø)
pymc/model/transform/optimization.py	97.63% <96.36%> (-2.37%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.