From 4ed93b1b53ec36c7444ef2009593bab3dab133d6 Mon Sep 17 00:00:00 2001
From: Adrian Seyboldt <adrian.seyboldt@gmail.com>
Date: Wed, 24 Jul 2024 19:27:44 +0200
Subject: [PATCH] feat: Support float32 settings in pytensor

---
 python/nutpie/compile_pymc.py |  2 +-
 tests/test_pymc.py            | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py
index 506a661..aaee456 100644
--- 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
index a4d3bf3..8f70699 100644
--- 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: