feat: Support float32 settings in pytensor

pymc-devs · Jul 26, 2024 · 4ed93b1 · 4ed93b1
1 parent dc261a0
commit 4ed93b1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py
@@ -483,7 +483,7 @@ def _make_functions(model, *, mode, compute_grad, join_expanded):
         variables = [joined[slice_val] for slice_val in zip(joined_slices)]
 
     replacements = {
-        model.rvs_to_values[var]: value.reshape(shape) if len(shape) != 1 else value
+        model.rvs_to_values[var]: value.reshape(shape).astype(var.dtype)
         for var, shape, value in zip(
             model.free_RVs,
             joined_shapes,

diff --git a/tests/test_pymc.py b/tests/test_pymc.py
@@ -23,6 +23,21 @@ def test_pymc_model(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@parameterize_backends
+def test_pymc_model_float32(backend, gradient_backend):
+    import pytensor
+
+    with pytensor.config.change_flags(floatX="float32"):
+        with pm.Model() as model:
+            pm.Normal("a")
+
+        compiled = nutpie.compile_pymc_model(
+            model, backend=backend, gradient_backend=gradient_backend
+        )
+        trace = nutpie.sample(compiled, chains=1)
+        trace.posterior.a  # noqa: B018
+
+
 @parameterize_backends
 def test_blocking(backend, gradient_backend):
     with pm.Model() as model: