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: